lib/utils/
llm.rs

1
2use crate::adapters::llm::OutputFormat;
3use crate::core::prompts::TunedPrompts;
4use serde_json::Value;
5use serde_json::json;
6use once_cell::sync::Lazy;
7use crate::core::learn::Instruction;
8use anyhow::{Result,Context};
9
10pub trait FromLLMResponse: Sized {
11  fn from_llm_response(response: String) -> Result<Self>;
12  fn output_format() -> OutputFormat;
13}
14
15impl FromLLMResponse for String {
16  fn from_llm_response(response: String) -> Result<Self> {
17      Ok(response)
18  }
19
20  fn output_format() -> OutputFormat {
21      OutputFormat::String
22  }
23}
24
25impl FromLLMResponse for Value {
26  fn from_llm_response(response: String) -> Result<Self> {
27      serde_json::from_str(&response).context("Failed to parse JSON response")
28  }
29
30  fn output_format() -> OutputFormat {
31      OutputFormat::Json
32  }
33}
34
35impl FromLLMResponse for TunedPrompts {
36  fn from_llm_response(response: String) -> Result<Self> {
37      let response: Value = serde_json::from_str(&response)
38          .context("Failed to parse JSON response")?;
39      let response = response.as_object()
40          .context("Response is not an object")?;
41      let prompts = response
42          .get("prompts")
43          .context("Missing 'prompts' field")?
44          .as_array()
45          .context("'prompts' field is not an array")?;
46      let prompts = prompts
47          .iter()
48          .map(|prompt| prompt.as_str().unwrap_or_default().to_string())
49          .collect();
50      Ok(TunedPrompts { prompts })
51  }
52
53  fn output_format() -> OutputFormat {
54      OutputFormat::StrictJson(TUNED_PROMPTS_RESPONSE_FORMAT.clone())
55  }
56}
57
58impl FromLLMResponse for Instruction {
59  fn from_llm_response(response: String) -> Result<Self> {
60      let response: Value = serde_json::from_str(&response)
61          .context("Failed to parse JSON response")?;
62      let instruction = response.get("instruction")
63          .context("Missing 'instruction' field")?
64          .as_str()
65          .context("'instruction' is not a string")?
66          .to_string();
67      let response_text = response.get("response")
68          .context("Missing 'response' field")?
69          .as_str()
70          .context("'response' is not a string")?
71          .to_string();
72      Ok(Instruction { instruction, response: response_text })
73  }
74
75  fn output_format() -> OutputFormat {
76      OutputFormat::StrictJson(INSTRUCTION_RESPONSE_FORMAT.clone())
77  }
78}
79
80pub static TUNED_PROMPTS_RESPONSE_FORMAT: Lazy<Value> = Lazy::new(|| {
81  json!({
82      "type": "json_schema",
83      "json_schema": {
84          "name": "prompts",
85          "strict": true,
86          "schema": {
87              "type": "object",
88              "properties": {
89                  "prompts": {
90                      "type": "array",
91                      "items": {
92                          "type": "string"
93                      }
94                  }
95              },
96              "required": ["prompts"],
97              "additionalProperties": false
98          }
99      }
100  })
101});
102
103pub static INSTRUCTION_RESPONSE_FORMAT: Lazy<Value> = Lazy::new(|| {
104  json!({
105      "type": "json_schema",
106      "json_schema": {
107          "name": "instruction",
108          "strict": true,
109          "schema": {
110              "type": "object",
111              "properties": {
112                  "instruction": {
113                      "type": "string"
114                  },
115                  "response": {
116                      "type": "string"
117                  }
118              },
119              "required": ["instruction", "response"],
120              "additionalProperties": false
121          }
122      }
123  })
124});