kproc 0.7.0

//! Encoders, from text/image to graph.
use kproc_llm::prelude::*;

use crate::{processor, processors::rag, Result};
use kproc_pmacros::Processor;

/// Text encoder, from text to graph
#[derive(Processor, Default)]
#[streams(rag::graph::TextEncoderInputs, rag::graph::EncoderOutputs)]
pub struct TextEncoder<T: LargeLanguageModel + Sync + Send>
{
  model: T,
}

impl<T: LargeLanguageModel + Sync + Send> TextEncoder<T>
{
  /// Create a new text encoder for the model
  pub fn new(model: T) -> Self
  {
    Self { model }
  }
}

const TEXT_ENCODER_SYSTEM: &str = r#"You are a data scientist working for a company that is building a graph database. Your task is to extract information from data and convert it into a graph database. Provide a set of Nodes in the form [ENTITY_ID, TYPE, PROPERTIES] and a set of relationships in the form [ENTITY_ID_1, RELATIONSHIP, ENTITY_ID_2, PROPERTIES]. It is important that the ENTITY_ID_1 and ENTITY_ID_2 exists as nodes with a matching ENTITY_ID. If you can't pair a relationship with a pair of nodes don't add it. When you find a node or relationship you want to add try to create a generic TYPE for it that  describes the entity you can also think of it as a label.

Example:
Data: Alice lawyer and is 25 years old and Bob is her roommate since 2001. Bob works as a journalist. Alice owns a the webpage www.alice.com and Bob owns the webpage www.bob.com.
Output:
{
    "nodes": [["alice", "Person", {"age": 25, "occupation": "lawyer", "name":"Alice"}], ["bob", "Person", {"occupation": "journalist", "name": "Bob"}], ["alice.com", "Webpage", {"url": "www.alice.com"}], ["bob.com", "Webpage", {"url": "www.bob.com"}]],
    "relationships": [["alice", "roommate", "bob", {"start": 2021}], ["alice", "owns", "alice.com", {}], ["bob", "owns", "bob.com", {}]]
}

Generate a graph for the following prompt."#;

impl<T: LargeLanguageModel + Sync + Send> processor::ValueProcessor for TextEncoder<T>
{
  #[allow(clippy::manual_async_fn)] // https://github.com/rust-lang/rust-clippy/issues/12664
  fn process_one(
    &self,
    inputs: rag::graph::TextEncoderInputs,
  ) -> impl std::future::Future<Output = Result<Self::Outputs>> + Send
  {
    async {
      let prompt = GenerationPrompt::prompt(inputs.prompt)
        .system(TEXT_ENCODER_SYSTEM.to_string())
        .format(Format::Json);

      let res = self.model.generate(prompt)?.await?;
      println!("{}", res);

      Ok(rag::graph::EncoderOutputs {
        graph: serde_json::from_str(&res)?,
      })
    }
  }
}

#[cfg(feature = "image")]
mod image_encoder
{
  use crate::processors::rag;

  use super::*;

  /// Encode an image into a graph.
  #[derive(Processor)]
  #[streams(rag::graph::ImageEncoderInputs, rag::graph::EncoderOutputs)]
  pub struct ImageEncoder<
    TImage: LargeLanguageModel + Sync + Send,
    TText: LargeLanguageModel + Sync + Send,
  > {
    text_model: TText,
    image_model: TImage,
  }

  impl<TImage: LargeLanguageModel + Sync + Send, TText: LargeLanguageModel + Sync + Send>
    ImageEncoder<TImage, TText>
  {
    /// Create a new image encoder, using two models. The image_model is used to convert
    /// the image into a string, describing the image. The text_model is used to convert
    /// the description into a graph.
    pub fn new(image_model: TImage, text_model: TText) -> Self
    {
      Self {
        image_model,
        text_model,
      }
    }
  }

  impl<TImage: LargeLanguageModel + Sync + Send, TText: LargeLanguageModel + Sync + Send>
    processor::ValueProcessor for ImageEncoder<TImage, TText>
  {
    #[allow(clippy::manual_async_fn)] // https://github.com/rust-lang/rust-clippy/issues/12664
    fn process_one(
      &self,
      inputs: Self::Inputs,
    ) -> impl std::future::Future<Output = Result<Self::Outputs>> + Send
    {
      async move {
        let image = inputs.prompt.read()?.to_owned();
        let res = self
          .image_model
          .generate(GenerationPrompt::prompt("Describe the given image.").image(image))?
          .await?;

        let prompt = GenerationPrompt::prompt(res)
          .system(TEXT_ENCODER_SYSTEM.to_string())
          .format(Format::Json);

        let res = self.text_model.generate(prompt)?.await?;
        Ok(rag::graph::EncoderOutputs {
          graph: serde_json::from_str(&res)?,
        })
      }
    }
  }
}

#[cfg(feature = "image")]
pub use image_encoder::ImageEncoder;