feat: support function cascade compress according to id (cloudwego#21)

* feat: support function cascade compress according to id

* feat: add debug logs

* feat: add compress for method

* feat: add compress all for functions/methods

Author: welkeyever

src/compress/compress.rs

@@ -13,13 +13,13 @@ pub struct Repository {
     #[serde(rename = "ModName")]
     mod_name: String,
     #[serde(rename = "Packages")]
-    packages: HashMap<String, Package>,
+    pub packages: HashMap<String, Package>,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Package {
     #[serde(rename = "Functions")]
-    functions: HashMap<String, Function>,
+    pub functions: HashMap<String, Function>,
     #[serde(rename = "Types")]
     types: HashMap<String, Struct>,
 }
@@ -44,6 +44,9 @@ pub struct Function {
     internal_method_calls: Option<HashMap<String, Identity>>,
     #[serde(rename = "ThirdPartyMethodCalls")]
     third_party_method_calls: Option<HashMap<String, Identity>>,
+
+    // compress_data
+    compress_data: Option<String>,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -62,14 +65,17 @@ pub struct Struct {
     inline_struct: Option<HashMap<String, Identity>>,
     #[serde(rename = "Methods")]
     methods: Option<HashMap<String, Identity>>,
+
+    // compress_data
+    compress_data: Option<String>,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Identity {
     #[serde(rename = "PkgPath")]
-    pkg_path: String,
+    pub pkg_path: String,
     #[serde(rename = "Name")]
-    name: String,
+    pub name: String,
 }
 
 pub fn from_json(json: &str) -> Result<Repository, Box<dyn Error>> {
@@ -78,49 +84,100 @@ pub fn from_json(json: &str) -> Result<Repository, Box<dyn Error>> {
 }
 
 
+pub async fn compress_all(repo: &mut Repository) {
+    let mut to_compress = 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)
+        }
+    }
+
+    for id in to_compress {
+        cascade_compress_function(&id, repo).await;
+    }
+}
+
 #[async_recursion]
-pub async fn cascade_compress_function(func: &mut Function) {
-    // if func.internal_function_calls.is_none() {
-    //     llm_compress(func);
-    //     return;
-    // }
-    //
-    // for (_, mut f) in func.internal_function_calls.as_mut().unwrap() {
-    //     if f.compress_info.is_none() {
-    //         cascade_compress_function(&mut f).await
-    //     }
-    // }
-    //
-    // llm_compress(func).await;
-    //
-    // return;
+pub async fn cascade_compress_function(id: &Identity, repo: &mut Repository) {
+    let mut to_compress = Vec::new();
+
+    {
+        let func_opt = repo.packages.get(id.pkg_path.as_str()).unwrap().functions.get(id.name.as_str());
+        if func_opt.is_none() {
+            println!("not found function, id {:?}", id);
+        }
+        let func_opt = func_opt.unwrap();
+        if func_opt.compress_data.is_some() {
+            println!("{} is already compressed, skip it.", func_opt.name);
+            return;
+        }
+        if let Some(calls) = &func_opt.internal_function_calls {
+            for (_, f) in calls {
+                let id = Identity { pkg_path: f.pkg_path.clone(), name: f.name.clone() };
+                to_compress.push(id);
+            }
+        }
+
+        if let Some(calls) = &func_opt.internal_method_calls {
+            for (_, f) in calls {
+                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_function(&f_id, repo).await;
+    }
+
+    let mut map = HashMap::new();
+    let content = {
+        let func_opt = repo.packages.get(id.pkg_path.as_str()).unwrap().functions.get(id.name.as_str()).unwrap();
+        if let Some(calls) = &func_opt.internal_function_calls {
+            for (k, f) in calls {
+                let sub_function = repo.packages.get(f.pkg_path.as_str()).unwrap().functions.get(f.name.as_str()).unwrap();
+                map.insert(k.clone(), sub_function.compress_data.clone().unwrap());
+            }
+        }
+
+        if let Some(calls) = &func_opt.internal_method_calls {
+            for (k, f) in calls {
+                let sub_function = repo.packages.get(f.pkg_path.as_str()).unwrap().functions.get(f.name.as_str()).unwrap();
+                map.insert(k.clone(), sub_function.compress_data.clone().unwrap());
+            }
+        }
+
+        println!("start to compress function: {}", func_opt.name);
+        llm_compress(func_opt.content.as_str(), map).await
+    };
+
+    let mut func_opt = repo.packages.get_mut(id.pkg_path.as_str()).unwrap().functions.get_mut(id.name.as_str()).unwrap();
+    func_opt.compress_data = content;
 }
 
-async fn llm_compress(func: &mut Function) {
-    // let mut map = HashMap::new(); // 创建一个空的 HashMap
-    // if func.internal_function_calls.is_some() {
-    //     for (k, ff) in func.internal_function_calls.as_ref().unwrap() {
-    //         map.insert(k.clone(), ff.compress_info.clone().unwrap());
-    //     }
-    // }
-    //
-    // let compress_data = _ollama_compress(func.content.clone(), map).await;
-    // func.compress_info = Option::from(compress_data);
+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)
 }
 
 pub async fn _ollama_compress(func: String, ctx: HashMap<String, String>) -> String {
     let request_url = format!("http://localhost:11434/api/generate");
 
-    let mut prompt = r#"You are an engineer who is proficient in Golang. You are responsible for summarizing the functions/methods given by the user.Try to condense output into one sentence and retain key information as much as possible. DO NOT show any codes in your answer. Function/methods content is as follow:"#.to_string();
+    let mut prompt = r#"You are an engineer who is proficient in Golang. You are responsible for summarizing the functions/methods given by the user, DO NOT put anything that is not mentioned in the function.Try to condense output into one sentence and retain key information as much as possible. DO NOT show any codes in your answer. Function/methods content is as follow:"#.to_string();
 
-    prompt.push_str("\n");
+    prompt.push_str("\n```\n");
     prompt.push_str(func.as_str());
-    prompt.push_str("\nRelated function:\n");
-    for (name, compressed_data) in ctx {
-        prompt.push_str(&*(name + ": " + &*compressed_data + "\n"));
+    prompt.push_str("\n```\n");
+    if !ctx.is_empty() {
+        prompt.push_str("\nRelated function:\n");
+        for (name, compressed_data) in ctx {
+            prompt.push_str(&*(name + ": " + &*compressed_data + "\n"));
+        }
     }
 
 
+    println!("use prompt:\n{}", prompt);
     let req_body: ollama_req = ollama_req { model: "codellama".to_string(), prompt };
     let client = reqwest::Client::new();
     let mut response = client