opt: (#7)

* feat: add module compress

* opt

Author: AsterDY

src/compress/compress.rs

@@ -31,6 +31,7 @@ use crate::config::CONFIG;
 use crate::storage::cache::get_cache;
 use crate::storage::cache::load_repo;
 
+use super::types::types::Module;
 use super::types::types::ToCompressVar;
 use super::types::types::Variant;
 
@@ -50,13 +51,17 @@ pub fn from_json(id: &str, json: &str) -> Result<Repository, Box<dyn Error>> {
     Ok(f)
 }
 
+pub fn is_externa_mod(modpath: &str) -> bool {
+    return modpath == "" || modpath.contains("@");
+}
+
 pub async fn compress_all(repo: &mut Repository) {
     let mut to_compress_func = Vec::new();
     let mut to_compress_type = Vec::new();
     let mut to_compress_var = Vec::new();
 
-    for (_, _mod) in &repo.modules {
-        if _mod.dir == "" {
+    for (mpath, _mod) in &repo.modules {
+        if is_externa_mod(mpath) {
             // NOTICE: empty dir means it's a external module, which is only used for lookup symbols
             continue;
         }
@@ -93,17 +98,32 @@ pub async fn compress_all(repo: &mut Repository) {
         cascade_compress_struct(&id, repo, &mut m).await;
     }
 
-    for (mname, _mod) in &repo.clone().modules {
-        if _mod.dir == "" {
+    for (mpath, _mod) in repo.clone().modules {
+        if is_externa_mod(&mpath) {
             // NOTICE: empty dir means it's a external module, which is only used for lookup symbols
             continue;
         }
         for (id, pkg) in &_mod.packages {
             if pkg.compress_data.is_none() {
-                compress_package(&id, mname, repo).await;
+                compress_package(&id, &mpath, repo).await;
             }
         }
+        compress_module(&mpath, repo).await;
+        repo.save_to_cache();
+    }
+}
+
+pub async fn compress_module(modpath: &str, repo: &mut Repository) {
+    let module = repo.modules.get_mut(modpath).unwrap();
+    let compress_data = module.to_compress();
+    let compress_data =
+        llm_compress_module(serde_json::to_string(&compress_data).unwrap().as_str()).await;
+    if compress_data.is_none() {
+        return;
     }
+    let compress_data = compress_data.unwrap();
+    module.compress_data = Some(compress_data);
+    println!("finish to compress module: {}", module.name);
 }
 
 pub async fn compress_package(id: &str, module: &str, repo: &mut Repository) {
@@ -811,6 +831,12 @@ pub async fn cascade_compress_struct(
     // panic!("empty compress for {:?}", id)
 }
 
+async fn llm_compress_module(m: &str) -> Option<String> {
+    let compressing = ToCompress::ToCompressModule(m.to_string());
+    let compress_data = compress(&compressing).await;
+    Option::from(compress_data)
+}
+
 async fn llm_compress_package(pkg: &str) -> Option<String> {
     let compress_pkg = ToCompress::ToCompressPkg(pkg.to_string());
     let compress_data = compress(&compress_pkg).await;

src/compress/llm/mod.rs

@@ -24,6 +24,7 @@ pub enum ToCompress {
     ToCompressType(String),
     ToCompressVar(String),
     ToCompressPkg(String),
+    ToCompressModule(String),
     ToMergeRustPkg(String),
     ToValidateRust(String),
 }

src/compress/llm/prompts.rs

@@ -30,6 +30,9 @@ macro_rules! choose_prompt_lang {
 
 pub fn make_compress_prompt(to_compress: &ToCompress) -> String {
     match to_compress {
+        ToCompress::ToCompressModule(f) => {
+            choose_prompt_lang!(PROMPT_COMPRESS_MOD).replace("{{DATA}}", f)
+        }
         ToCompress::ToCompressType(t) => {
             choose_prompt_lang!(PROMPT_COMPRESS_TYPE).replace("{{DATA}}", t)
         }
@@ -462,7 +465,7 @@ bufferSizeLimit is an integer variable with an initial value of integer 1024 tha
 
 const PROMPT_COMPRESS_PKG_ZH: &str = r##"
 # 角色
-你是一名经验丰富的工程师，专门研究lang，并深入了解其各种包。你的主要职责是利用其他开发人员提供的有关公共函数和类型的数据，简化并总结lang包的基本功能。你的目标是通过使这些包更易理解，帮助对这些包了解较少的工程师。
+你是一名经验丰富的工程师，并深入了解其各种包。你的主要职责是利用其他开发人员提供的有关公共函数和类型的数据，简化并总结包的基本功能。你的目标是通过使这些包更易理解，帮助对这些包了解较少的工程师。
 
 # 提示
 
@@ -655,3 +658,124 @@ Key global variables:
 {{DATA}}
 
 "##;
+
+const PROMPT_COMPRESS_MOD_ZH: &str = r##"
+# 角色
+你是一名经验丰富的工程师，并深入了解其各种模块。你的主要职责是利用其他开发人员提供的有关公共函数和类型的数据，简化并总结模块的基本功能。你的目标是通过使这些包更易理解，帮助对这些包了解较少的工程师。
+
+# 提示
+
+## 输入格式(JSON)
+包含 一个具体的类型定义 及 其依赖的其他语言符号描述：
+- "Name": 模块名称
+- "Dir": 模块所处的相对仓库位置
+    - "Description": 该方法的代码或总结, 格式为字符串
+- "Packages": 格式为数组。该数组中的每个对象表示此模块内每个子包描述：
+    - "Name": 使用该包的名称，格式为字符串
+    - "Description": 该包的总结，格式为字符串
+
+
+## 输出格式（text）
+直接输出总结内容。不要输出JSON（IMPORTANT）！
+
+## 总结内容
+- 该模块的主要功能和用途
+- 该模块的一些关键函数和类型的描述
+
+
+# 约束
+- 专注于高度总结模块的基本功能，避免深入具体实现细节。
+- 编写简短且易于理解的总结，供其他工程师参考。
+- 保持与提供的输入数据一致的技术术语。
+- 输出字符限制为2000字符。
+
+
+# 示例
+
+## 输入
+{
+    "Name": "github.com/cloudwego/localsession",
+    "Dir": ".",
+    "Packages": [
+        {
+            "Description": "该包用于管理会话上下文，并定义了通用的Session接口",
+            "Name": "github.com/cloudwego/localsession",
+        },
+        {
+            "Description": "该包用于处理具体的上下文的metainfo等信息的兜底方式",
+            "Name": "github.com/cloudwego/localsession/backup",
+        }
+    ]
+}
+
+## 输出
+此模块位于当前目录，为应用中的会话管理提供工具，特别关注会话上下文的备份和恢复机制。它包括创建、备份、恢复和清除会话上下文的功能，以及默认的上下文兜底方式。
+关键包：
+    - github.com/cloudwego/localsession: 该包用于管理会话上下文，并定义了通用的Session接口
+    - github.com/cloudwego/localsession/backup: 该包用于处理具体的上下文的metainfo等信息的兜底方式
+
+
+# 现在，请开始处理如下输入：
+
+{{DATA}}
+
+"##;
+
+const PROMPT_COMPRESS_MOD_EN: &str = r##"
+# Character
+You are an experienced engineer and have in-depth knowledge of its various modules. Your primary responsibility is to simplify and summarize the basic functionality of the module using data about common functions and types provided by other developers. Your goal is to help engineers who know less about these packages by making them easier to understand.
+
+# Tips
+
+## Input format (JSON)
+Contains a specific type definition and descriptions of other language symbols on which it depends:
+- Name: indicates the module name
+- "Dir": indicates the relative repository location of the module
+- "Description": indicates the code or summary of the method. The format is a string
+- "Packages": in array format. Each object in this array represents each subpackage description within this module:
+- "Name": indicates the name of the package. The format is a string
+- "Description": indicates the summary of the package. The format is a string
+
+
+## Output format (text)
+Output summary content directly. Do not output JSON (IMPORTANT)!
+
+## Summarize the content
+- Main functions and uses of the module
+- Description of some of the key functions and types of the module
+
+
+# Constraint
+- Focus on the basic functions of highly summarized modules and avoid delving into specific implementation details.
+- Write short and easy to understand summaries for other engineers to refer to.
+- Technical terms that are consistent with the input data provided.
+- The output character limit is 2000 characters.
+
+
+# Examples
+
+## Input
+{
+    "Name": "github.com/cloudwego/localsession",
+    "Dir": ".",
+    "Packages": [
+            {
+                "Description": "This package manages the Session context and defines the generic session interface ",
+                "Name": "github.com/cloudwego/localsession",
+            },
+            {
+                "Description": "This package is used to handle the specific context of metainfo and other information in the back of the way ",
+                "Name": "github.com/cloudwego/localsession/backup",
+            }
+        ]
+}
+
+## Output
+This module is located in the current directory and provides tools for session management in the application, with a special focus on backup and recovery mechanisms for session context. It includes the ability to create, back up, restore, and clear session context, as well as the default context bypass.
+Key package:
+- github.com/cloudwego/localsession: this package is used to manage Session context, and defines the general Session interface
+- github.com/cloudwego/localsession/backup: this package to deal with the specific context of the information such as the metainfo way out
+
+# Now, please summarize below input：
+{{DATA}}
+"##;

src/compress/types/types.rs

@@ -239,6 +239,7 @@ impl Repository {
                 self.modules.insert(mod_name.clone(), _mod.clone());
             }
         }
+        self.graph = other.graph.clone();
     }
 
     pub fn save_to_cache(&self) {
@@ -357,9 +358,29 @@ pub struct Module {
     #[serde(rename = "Packages")]
     pub packages: HashMap<String, Package>,
     #[serde(rename = "Files")]
-    pub files: HashMap<String, File>,
-    #[serde(rename = "Language")]
+    pub files: Option<HashMap<String, File>>,
+    #[serde(rename = "Language", default)]
     pub language: String,
+    #[serde(rename = "compress_data")]
+    pub compress_data: Option<String>,
+}
+
+impl Module {
+    pub fn to_compress(&self) -> ToCompressModule {
+        let mut packages = Vec::new();
+        for (_, p) in self.packages.iter() {
+            packages.push(Description {
+                name: &p.id,
+                description: p.compress_data.as_ref().unwrap(),
+            });
+        }
+
+        ToCompressModule {
+            name: &self.name,
+            dir: &self.dir,
+            packages: Some(packages),
+        }
+    }
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -769,6 +790,16 @@ pub(crate) struct ToCompressPkg<'a> {
     pub(crate) vars: Option<Vec<Description<'a>>>,
 }
 
+#[derive(Serialize, Debug)]
+pub(crate) struct ToCompressModule<'a> {
+    #[serde(rename = "Name")]
+    pub(crate) name: &'a str,
+    #[serde(rename = "Dir")]
+    pub(crate) dir: &'a str,
+    #[serde(rename = "Packages")]
+    pub(crate) packages: Option<Vec<Description<'a>>>,
+}
+
 #[derive(Serialize, Debug)]
 pub(crate) struct Description<'a> {
     pub name: &'a str,

src/lang/patch/lib.go

@@ -55,9 +55,10 @@ type Options struct {
 	OutDir  string
 }
 
-func NewPatcher(opts Options) *Patcher {
+func NewPatcher(repo *uniast.Repository, opts Options) *Patcher {
 	return &Patcher{
 		Options: opts,
+		repo:    repo,
 	}
 }
 
@@ -98,6 +99,7 @@ next_dep:
 		})
 	}
 	n := patchNode{
+		Identity: patch.Id,
 		FileLine: node.FileLine(),
 		Codes:    patch.Codes,
 		File:     f,
@@ -174,7 +176,7 @@ func (p *Patcher) Flush() error {
 	}
 
 	// write origins
-	for _, mod := range p.repo.Modules {
+	for _, mod := range p.repo.InternalModules() {
 		for _, f := range mod.Files {
 			if p.patches[f.Path] != nil {
 				continue

src/lang/uniast/ast.go

@@ -37,6 +37,16 @@ type Repository struct {
 	Graph   map[string]*Node
 }
 
+func (r Repository) InternalModules() []*Module {
+	var ret []*Module
+	for k, v := range r.Modules {
+		if !IsExternalModule(k) {
+			ret = append(ret, v)
+		}
+	}
+	return ret
+}
+
 func NewRepository(name string) Repository {
 	ret := Repository{
 		Name:    name,
@@ -82,6 +92,7 @@ type Module struct {
 	Packages     map[PkgPath]*Package // pkage import path => Package
 	Dependencies map[string]string    // module name => module_path@version
 	Files        map[string]*File     // relative path => file info
+	CompressData *string              `json:"compress_data,omitempty"` // module compress info
 }
 
 func (r Repository) GetFileById(id Identity) *File {
@@ -164,9 +175,10 @@ func ModPathName(mod ModPath) string {
 
 // Identity holds identity information about a third party declaration
 type Identity struct {
-	ModPath        // ModPath is the module which the package belongs to
-	PkgPath        // Import Path of the third party package
-	Name    string // Unique Name of declaration (FunctionName, TypeName.MethodName, InterfaceName<TypeName>.MethodName, or TypeName)
+	ModPath `json:"ModPath" jsonschema:"description=the compiling module of the ast node, the format is {ModName} or {ModName}@{Version}"` // ModPath is the module which the package belongs to
+	PkgPath `json:"PkgPath" jsonschema:"description=the namespace of the ast node"`                                                        // Import Path of the third party package
+
+	Name string `json:"Name" jsonschema:"description=unique name of the ast node, the format is one of {FunctionName}, {TypeName}.{MethodName}, {InterfaceName}<{TypeName}>.{MethodName}, {TypeName}"` // Unique Name of declaration (FunctionName, TypeName.MethodName, InterfaceName<TypeName>.MethodName, or TypeName)
 }
 
 func NewIdentity(mod, pkg, name string) Identity {

src/lang/uniast/node.go

@@ -14,10 +14,6 @@
 
 package uniast
 
-import (
-	"fmt"
-)
-
 func (r *Repository) GetNode(id Identity) *Node {
 	key := id.Full()
 	node, ok := r.Graph[key]
@@ -55,6 +51,7 @@ func (r *Repository) SetNode(id Identity, typ NodeType) *Node {
 		node = &Node{
 			Identity: id,
 			Type:     typ,
+			Repo:     r,
 		}
 		r.Graph[key] = node
 	}
@@ -89,6 +86,7 @@ func (r *Repository) AddRelation(node *Node, dep Identity) {
 	if !ok {
 		nd = &Node{
 			Identity: dep,
+			Repo:     r,
 		}
 		r.Graph[key] = nd
 	}
@@ -209,9 +207,6 @@ func (t NodeType) MarshalJSON() ([]byte, error) {
 
 func (t *NodeType) UnmarshalJSON(b []byte) error {
 	typ := NewNodeType(string(b))
-	if typ == UNKNOWN {
-		return fmt.Errorf("unknown node type: %s", b)
-	}
 	*t = typ
 	return nil
 }