retrieval-pipeline: v0

retrieval-pipeline/src/answer_retrieval.rs

@@ -0,0 +1,117 @@
+use async_openai::{
+    error::OpenAIError,
+    types::{
+        ChatCompletionRequestSystemMessage, ChatCompletionRequestUserMessage,
+        CreateChatCompletionRequest, CreateChatCompletionRequestArgs, CreateChatCompletionResponse,
+        ResponseFormat, ResponseFormatJsonSchema,
+    },
+};
+use common::{
+    error::AppError,
+    storage::types::{
+        message::{format_history, Message},
+        system_settings::SystemSettings,
+    },
+};
+use serde::Deserialize;
+use serde_json::Value;
+
+use super::answer_retrieval_helper::get_query_response_schema;
+
+#[derive(Debug, Deserialize)]
+pub struct Reference {
+    #[allow(dead_code)]
+    pub reference: String,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct LLMResponseFormat {
+    pub answer: String,
+    #[allow(dead_code)]
+    pub references: Vec<Reference>,
+}
+
+#[derive(Debug)]
+pub struct Answer {
+    pub content: String,
+    pub references: Vec<String>,
+}
+
+pub fn create_user_message(entities_json: &Value, query: &str) -> String {
+    format!(
+        r"
+        Context Information:
+        ==================
+        {entities_json}
+
+        User Question:
+        ==================
+        {query}
+        "
+    )
+}
+
+pub fn create_user_message_with_history(
+    entities_json: &Value,
+    history: &[Message],
+    query: &str,
+) -> String {
+    format!(
+        r"
+        Chat history:
+        ==================
+        {}
+        
+        Context Information:
+        ==================
+        {}
+
+        User Question:
+        ==================
+        {}
+        ",
+        format_history(history),
+        entities_json,
+        query
+    )
+}
+
+pub fn create_chat_request(
+    user_message: String,
+    settings: &SystemSettings,
+) -> Result<CreateChatCompletionRequest, OpenAIError> {
+    let response_format = ResponseFormat::JsonSchema {
+        json_schema: ResponseFormatJsonSchema {
+            description: Some("Query answering AI".into()),
+            name: "query_answering_with_uuids".into(),
+            schema: Some(get_query_response_schema()),
+            strict: Some(true),
+        },
+    };
+
+    CreateChatCompletionRequestArgs::default()
+        .model(&settings.query_model)
+        .messages([
+            ChatCompletionRequestSystemMessage::from(settings.query_system_prompt.clone()).into(),
+            ChatCompletionRequestUserMessage::from(user_message).into(),
+        ])
+        .response_format(response_format)
+        .build()
+}
+
+pub async fn process_llm_response(
+    response: CreateChatCompletionResponse,
+) -> Result<LLMResponseFormat, AppError> {
+    response
+        .choices
+        .first()
+        .and_then(|choice| choice.message.content.as_ref())
+        .ok_or(AppError::LLMParsing(
+            "No content found in LLM response".into(),
+        ))
+        .and_then(|content| {
+            serde_json::from_str::<LLMResponseFormat>(content).map_err(|e| {
+                AppError::LLMParsing(format!("Failed to parse LLM response into analysis: {e}"))
+            })
+        })
+}