feat: struct compress (cloudwego#23)

* feat: add model file

* optimzie: prompt engineering

* feat: refactor & support type compress

* fix: errors

* refactor: fix llm related error

* chore: remove unused import

* feat: prompt engineering & change to mistral model

Author: welkeyever

script/ModelFile

@@ -0,0 +1,28 @@
+FROM mistral:latest
+
+SYSTEM """# Character
+          You're a Golang engineer with a special talent for condensing and summarizing complex methods or functions. You've a knack for retaining key information while keeping the summary concise and informative. Your task involves summarizing information provided to you through JSON strings, although you never show any code in your summaries or refer to the JSON string format.
+
+          ## Skills
+          ### Skill 1: Summarize Method or Function
+          - Extract the content of the main function/method to be summarized from the "Content" field.
+          - Closely look at the "Related_func" list for all the functions or methods that are called in the main function/method.
+          - Be attentive to the possibilities that "Related_func" can be null which means there is no other function/method is called. Just ignore "Related_func" in this case. And do not mention the lack of "Related_func" in your summaries.
+          - Your summary should decode the functionality of the method or function in a concise sentence without losing its essence.
+          - Stick to the format: "XXX is for XXX..."
+
+          ### Skill 2: Consider Related Functions or Methods
+          - Each element from the "Related_func" list requires thorough consideration. Notably, the "CallName" shows the name that is called in the main function/method while "Description" gives the summarized context of it.
+          - Integrate the information from "Related_func"(if there has any) to your summary of the main method or function, enriching it with more specific detail.
+
+          ## Constraints:
+          - Remember to keep your summaries focused on the function/method, ignore any mentions of "JSON string".
+          - Never show any code in your summaries.
+          - Maintain the output format by beginning the summarization as 'XXX is for XXX...'
+          - Remain on the topic of summarizing functions or methods, if the user diverges from this, do not accommodate their queries.
+          - Match the language used by the user in their queries. Make sure not to use a language that doesn't parallel the user's choice.
+          - Only answer questions regarding function or method summarization. For any unrelated queries, do not provide an answer."""
+
+PARAMETER temperature 0.7
+
+

script/ModelFile-func

@@ -0,0 +1,30 @@
+FROM mistral:latest
+
+SYSTEM """# Character
+          You're a systematic Go programmer. You specialize in summarizing key details of user-provided Go types (such as structs, interfaces, and so on), which are pitched in JSON format.
+
+          ## Skills
+
+          ### Skill 1: Summarize Type Content
+          - Inspect and recognize the definition represented in the "Content" of the type in JSON format.
+
+          ### Skill 2: Identify Related Methods
+          - Evaluate the "Related_methods" list, each item of which is an object representing a related method for this type.
+          - Every object contains "Name" as the method's name and "Description" as the condensed information of the method.
+          - If no methods exist for this type, "Related_methods" will be null.
+
+          ### Skill 3: Identify Related Types
+          - Comprehend the "Related_types" list, each element of which is an object illustrating a type that is utilized in the type definition.
+          - Each object encompasses "Name" as the type's name and "Description" as the brief information of the type.
+          - If there are no types for this type, "Related_types" will be null.
+
+          ## Constraints:
+          - Remember to keep your summaries focused on the type, ignore any mentions of "JSON string".
+          - Never show any code in your summaries.
+          - It is strictly prohibited to show the origin JSON string or any part of it.
+          - Using the information provided in JSON format, you should accurately summarize the type content, related methods, and related types.
+          - "Related_methods" and "Related_types", should be null when the type doesn't have any methods or other types associated.
+          - Start your responses with a summary of the type content directly."""
+
+
+PARAMETER temperature 0.7

script/ModelFile-type

@@ -82,37 +82,66 @@ pub struct Identity {
 struct ToCompressFunc {
     #[serde(rename = "Content")]
     content: String,
-    #[serde(rename = "related_func")]
-    related_func: Option<Vec<CalledFunc>>,
+    #[serde(rename = "Related_func")]
+    related_func: Option<Vec<CalledType>>,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+struct ToCompressType {
+    #[serde(rename = "Content")]
+    content: String,
+    #[serde(rename = "Related_methods")]
+    related_methods: Option<Vec<KeyValueType>>,
+    #[serde(rename = "Related_types")]
+    related_types: Option<Vec<KeyValueType>>,
 }
 
 
 #[derive(Serialize, Deserialize, Debug)]
-struct CalledFunc {
+struct CalledType {
     #[serde(rename = "CallName")]
     pub call_name: String,
     #[serde(rename = "Description")]
     pub description: String,
 }
 
+#[derive(Serialize, Deserialize, Debug)]
+struct KeyValueType {
+    #[serde(rename = "Name")]
+    pub name: String,
+    #[serde(rename = "Description")]
+    pub description: String,
+}
+
 pub fn from_json(json: &str) -> Result<Repository, Box<dyn Error>> {
     let f: Repository = serde_json::from_str(json)?;
     Ok(f)
 }
 
 
 pub async fn compress_all(repo: &mut Repository) {
-    let mut to_compress = Vec::new();
+    let mut to_compress_func = Vec::new();
+    let mut to_compress_type = Vec::new();
+
     for (_, pkg) in &repo.packages {
         for (_, func) in &pkg.functions {
             let id = Identity { pkg_path: func.pkg_path.clone(), name: func.name.clone() };
-            to_compress.push(id)
+            to_compress_func.push(id)
+        }
+
+        for (_, _type) in &pkg.types {
+            let id = Identity { pkg_path: _type.pkg_path.clone(), name: _type.name.clone() };
+            to_compress_type.push(id)
         }
     }
 
-    for id in to_compress {
+    for id in to_compress_func {
         cascade_compress_function(&id, repo).await;
     }
+
+    for id in to_compress_type {
+        cascade_compress_struct(&id, repo).await;
+    }
 }
 
 #[async_recursion]
@@ -170,7 +199,7 @@ pub async fn cascade_compress_function(id: &Identity, repo: &mut Repository) {
             println!("content is empty skip it");
             Some("".to_string())
         } else {
-            llm_compress(func_opt.content.as_str(), map).await
+            llm_compress_func(func_opt.content.as_str(), map).await
         }
     };
 
@@ -183,30 +212,175 @@ pub async fn cascade_compress_function(id: &Identity, repo: &mut Repository) {
     func_opt.compress_data = content;
 }
 
-async fn llm_compress(func: &str, extra: HashMap<String, String>) -> Option<String> {
-    let compress_data = _ollama_compress(func.to_string(), extra).await;
-    Option::from(compress_data)
+#[async_recursion]
+pub async fn cascade_compress_struct(id: &Identity, repo: &mut Repository) {
+    let mut to_compress = Vec::new();
+
+    {
+        let struct_opt = repo.packages.get(id.pkg_path.as_str()).unwrap().types.get(id.name.as_str());
+        if struct_opt.is_none() {
+            println!("not found struct, id {:?}", id);
+        }
+        let stru = struct_opt.unwrap();
+        if stru.compress_data.is_some() {
+            println!("{} is already compressed, skip it.", stru.name);
+            return;
+        }
+        if let Some(sub) = &stru.sub_struct {
+            for (_, f) in sub {
+                if !f.pkg_path.starts_with(&repo.mod_name) {
+                    continue;
+                }
+
+                let id = Identity { pkg_path: f.pkg_path.clone(), name: f.name.clone() };
+                to_compress.push(id);
+            }
+        }
+
+        if let Some(inline) = &stru.inline_struct {
+            for (_, f) in inline {
+                if !f.pkg_path.starts_with(&repo.mod_name) {
+                    continue;
+                }
+
+                let id = Identity { pkg_path: f.pkg_path.clone(), name: f.name.clone() };
+                to_compress.push(id);
+            }
+        }
+    }
+
+    for f_id in to_compress {
+        cascade_compress_struct(&f_id, repo).await;
+    }
+
+    let mut type_map = HashMap::new();
+    let mut method_map = HashMap::new();
+    let content = {
+        let _type = repo.packages.get(id.pkg_path.as_str()).unwrap().types.get(id.name.as_str()).unwrap();
+        if let Some(subs) = &_type.sub_struct {
+            for (k, f) in subs {
+                let pkg = repo.packages.get(f.pkg_path.as_str());
+                if pkg.is_none() {
+                    // TODO
+                    eprintln!("do not get the type, must be a third party one: {:?}", f);
+                    continue;
+                }
+                let sub = pkg.unwrap().types.get(f.name.as_str());
+                type_map.insert(k.clone(), sub.unwrap().compress_data.clone().unwrap());
+            }
+        }
+
+        if let Some(inlines) = &_type.inline_struct {
+            for (k, f) in inlines {
+                let pkg = repo.packages.get(f.pkg_path.as_str());
+                if pkg.is_none() {
+                    // TODO
+                    eprintln!("do not get the type, must be a third party one: {:?}", f);
+                    continue;
+                }
+                let inline = repo.packages.get(f.pkg_path.as_str()).unwrap().types.get(f.name.as_str());
+
+                type_map.insert(k.clone(), inline.unwrap().compress_data.clone().unwrap());
+            }
+        }
+
+        if let Some(methods) = &_type.methods {
+            for (k, f) in methods {
+                let func = repo.packages.get(f.pkg_path.as_str()).unwrap().functions.get(f.name.as_str());
+                if func.is_none() {
+                    // TODO
+                    eprintln!("[BUG] do not get the method of the type, id: {:?}", f);
+                } else {
+                    method_map.insert(k.clone(), func.unwrap().compress_data.clone().unwrap());
+                }
+            }
+        }
+
+        println!("start to compress type: {}", _type.name);
+        if _type.content.is_empty() {
+            println!("content is empty skip it");
+            Some("".to_string())
+        } else {
+            llm_compress_type(_type.content.as_str(), type_map, method_map).await
+        }
+    };
+
+    let mut type_opt = repo.packages.get_mut(id.pkg_path.as_str()).unwrap().types.get_mut(id.name.as_str()).unwrap();
+    if content.is_some() {
+        let content = content.unwrap().trim().to_string();
+        type_opt.compress_data = Some(content);
+        return;
+    }
+    type_opt.compress_data = content;
 }
 
 
-pub async fn _ollama_compress(func: String, ctx: HashMap<String, String>) -> String {
-    let request_url = format!("http://localhost:11434/api/generate");
+pub enum ToCompress {
+    ToCompressFunc(String),
+    ToCompressType(String),
+}
 
-    let mut compress_func = ToCompressFunc { content: func, related_func: None };
-    if !ctx.is_empty() {
+async fn llm_compress_func(func: &str, extra: HashMap<String, String>) -> Option<String> {
+    let mut compress_func = ToCompressFunc { content: func.to_string(), related_func: None };
+    if !extra.is_empty() {
         let mut related_func = Vec::new();
-        for (name, compressed_data) in ctx {
-            let re = CalledFunc { call_name: name, description: compressed_data };
+        for (name, compressed_data) in extra {
+            let re = CalledType { call_name: name, description: compressed_data };
             related_func.push(re);
         }
         compress_func.related_func = Some(related_func);
     }
+    let to_compress_str = serde_json::to_string(&compress_func).unwrap();
+    let compress_func_enum = ToCompress::ToCompressFunc(to_compress_str);
+    let compress_data = _ollama_compress(compress_func_enum).await;
+    Option::from(compress_data)
+}
+
+// depends on the compressed info of methods, so call llm_compress_func first.
+async fn llm_compress_type(func: &str, extra_type: HashMap<String, String>, related_methods: HashMap<String, String>) -> Option<String> {
+    let mut compress_type = ToCompressType { content: func.to_string(), related_methods: None, related_types: None };
+    if !extra_type.is_empty() {
+        let mut r_type = Vec::new();
+        for (name, compressed_data) in extra_type {
+            let re = KeyValueType { name, description: compressed_data };
+            r_type.push(re);
+        }
+        compress_type.related_types = Some(r_type);
+    }
+
+    if !related_methods.is_empty() {
+        let mut r_methods = Vec::new();
+        for (name, compressed_data) in related_methods {
+            let re = KeyValueType { name, description: compressed_data };
+            r_methods.push(re);
+        }
+        compress_type.related_methods = Some(r_methods);
+    }
+
+
+    let to_compress_str = serde_json::to_string(&compress_type).unwrap();
+    let compress_type_enum = ToCompress::ToCompressType(to_compress_str);
+    let compress_data = _ollama_compress(compress_type_enum).await;
+    Option::from(compress_data)
+}
 
-    let to_compress_func = serde_json::to_string(&compress_func).unwrap();
 
+pub async fn _ollama_compress(to_compress: ToCompress) -> String {
+    let request_url = format!("http://localhost:11434/api/generate");
+    let mut model_name = "codellama-private";
+    let mut to_compress_str = String::new();
+    match to_compress {
+        ToCompress::ToCompressType(t) => {
+            model_name = "codellama-private-type";
+            to_compress_str = t;
+        }
+        ToCompress::ToCompressFunc(f) => {
+            to_compress_str = f;
+        }
+    }
 
-    println!("use prompt:\n{}", to_compress_func);
-    let req_body: ollama_req = ollama_req { model: "codellama-private".to_string(), prompt: to_compress_func };
+    println!("use prompt:\n{}", to_compress_str);
+    let req_body: OllamaReq = OllamaReq { model: model_name.to_string(), prompt: to_compress_str };
     let client = reqwest::Client::new();
     let mut response = client
         .post(&request_url)
@@ -222,7 +396,7 @@ pub async fn _ollama_compress(func: String, ctx: HashMap<String, String>) -> Str
             break;
         }
 
-        let value: ollama_resp = result.unwrap();
+        let value: OllamaResp = result.unwrap();
 
         if !value.response.is_empty() {
             output.push_str(value.response.as_str());
@@ -237,13 +411,13 @@ pub async fn _ollama_compress(func: String, ctx: HashMap<String, String>) -> Str
 }
 
 #[derive(Serialize, Deserialize, Debug)]
-struct ollama_req {
+struct OllamaReq {
     model: String,
     prompt: String,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
-struct ollama_resp {
+struct OllamaResp {
     model: String,
     created_at: String,
     response: String,