| 5 | import re |
| 6 | |
| 7 | class CodeMemory(Memory): |
| 8 | def __init__(self, project_path: str, db_name: str = '.sa', platform: str = 'OpenAI', api_key: str = None, embedding_model: str = "text-embedding-ada-002"): |
| 9 | super().__init__(project_path, db_name, platform, api_key, embedding_model) |
| 10 | self.collection_name = 'code_memory' |
| 11 | |
| 12 | def add_code_files(self, directory: str, exclude_prefix: List[str] = ["workplace_"]): |
| 13 | """ |
| 14 | Add all code files in the specified directory to the memory. |
| 15 | |
| 16 | Args: |
| 17 | directory (str): The directory path containing the code files to add. |
| 18 | """ |
| 19 | code_files = [] |
| 20 | for root, _, files in os.walk(directory): |
| 21 | root_name = str(root) |
| 22 | if any(prefix in root_name for prefix in exclude_prefix): |
| 23 | continue |
| 24 | for file in files: |
| 25 | |
| 26 | if file.endswith(('.py', '.js', '.java', '.cpp', '.h', '.c', '.html', '.css')): # add more file types if needed |
| 27 | file_path = os.path.join(root, file) |
| 28 | with open(file_path, 'r', encoding='utf-8') as f: |
| 29 | content = f.read() |
| 30 | code_files.append({ |
| 31 | "query": f"File: {file_path}\n\nContent:\n{content}", |
| 32 | "response": f"This is the content of file {file_path}" |
| 33 | }) |
| 34 | self.add_query(code_files, self.collection_name) |
| 35 | |
| 36 | def query_code(self, query_text: str, n_results: int = 5) -> List[Dict]: |
| 37 | """ |
| 38 | Query the code memory. |
| 39 | |
| 40 | Args: |
| 41 | query_text (str): The query text |
| 42 | n_results (int): The number of results to return |
| 43 | |
| 44 | Returns: |
| 45 | List[Dict]: The query results list |
| 46 | """ |
| 47 | results = self.query([query_text], self.collection_name, n_results) |
| 48 | return [ |
| 49 | { |
| 50 | "file": doc.split('\n')[0].replace("File: ", ""), |
| 51 | "content": '\n'.join(doc.split('\n')[3:]), |
| 52 | "metadata": metadata |
| 53 | } |
| 54 | for doc, metadata in zip(results['documents'][0], results['metadatas'][0]) |
| 55 | ] |
| 56 | |
| 57 | class CodeReranker(Reranker): |
| 58 | def __init__(self, model: str) -> None: |