| 9 | from openai import OpenAI |
| 10 | import pandas as pd |
| 11 | class CodeTreeMemory(Memory): |
| 12 | def __init__(self, project_path: str, db_name: str = '.code_tree', platform: str = 'OpenAI', api_key: str = None, embedding_model: str = "text-embedding-ada-002"): |
| 13 | super().__init__(project_path, db_name, platform, api_key, embedding_model) |
| 14 | self.collection_name = 'code_tree_memory' |
| 15 | self.embedder = OpenAI(api_key=api_key) |
| 16 | |
| 17 | |
| 18 | def add_code_files(self, directory: str, exclude_prefix: List[str] = ["workplace_"]): |
| 19 | """ |
| 20 | 将指定目录下的所有代码文件添加到内存中。 |
| 21 | |
| 22 | Args: |
| 23 | directory (str): 要添加的代码文件所在的目录路径 |
| 24 | """ |
| 25 | tree_sitter_parent_dir = os.path.dirname(os.getcwd()) |
| 26 | # Build Tree sitter Parser object |
| 27 | Language.build_library( |
| 28 | f"{tree_sitter_parent_dir}/my-languages.so", |
| 29 | [ |
| 30 | f"{tree_sitter_parent_dir}/tree-sitter-python", |
| 31 | ], |
| 32 | ) |
| 33 | parser = CodeParser( |
| 34 | language="python", |
| 35 | node_types=["class_definition", "function_definition"], |
| 36 | path_to_object_file=tree_sitter_parent_dir, |
| 37 | ) |
| 38 | logger.info("Parsing Code...") |
| 39 | parsed_snippets = parser.parse_directory( |
| 40 | directory |
| 41 | ) |
| 42 | snippet_texts = list(map(lambda x: x.snippet.decode("ISO-8859-1"), parsed_snippets)) |
| 43 | embedded_texts = self.embedder.embeddings.create(input=snippet_texts, model="text-embedding-3-small").data |
| 44 | embedded_snippets = [] |
| 45 | for code_text, embedding, snippet in zip( |
| 46 | snippet_texts, embedded_texts, parsed_snippets |
| 47 | ): |
| 48 | snippet.snippet = code_text |
| 49 | snippet.embedding = embedding.embedding |
| 50 | embedded_snippets.append(snippet) |
| 51 | |
| 52 | # Convert Snippets to DataFrame for ChromaDB Ingestion |
| 53 | data = pd.DataFrame(to_dataframe_row(embedded_snippets)) |
| 54 | collection = self.client.get_or_create_collection( |
| 55 | name=self.collection_name, metadata={"hnsw:space": "cosine"} |
| 56 | ) |
| 57 | logger.info( |
| 58 | f"Adding {data.shape[0]} Code snippets and embedding to " |
| 59 | "local chroma db collection..." |
| 60 | ) |
| 61 | collection.add( |
| 62 | documents=data["snippets"].tolist(), |
| 63 | embeddings=data["embeddings"].tolist(), |
| 64 | metadatas=data["metadatas"].tolist(), |
| 65 | ids=data["ids"].tolist(), |
| 66 | ) |
| 67 | |
| 68 | |