chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,174 @@
|
||||
import os
|
||||
from typing import List, Dict
|
||||
from autoagent.memory.rag_memory import Memory, Reranker
|
||||
from litellm import completion
|
||||
import re
|
||||
|
||||
class CodeMemory(Memory):
|
||||
def __init__(self, project_path: str, db_name: str = '.sa', platform: str = 'OpenAI', api_key: str = None, embedding_model: str = "text-embedding-ada-002"):
|
||||
super().__init__(project_path, db_name, platform, api_key, embedding_model)
|
||||
self.collection_name = 'code_memory'
|
||||
|
||||
def add_code_files(self, directory: str, exclude_prefix: List[str] = ["workplace_"]):
|
||||
"""
|
||||
Add all code files in the specified directory to the memory.
|
||||
|
||||
Args:
|
||||
directory (str): The directory path containing the code files to add.
|
||||
"""
|
||||
code_files = []
|
||||
for root, _, files in os.walk(directory):
|
||||
root_name = str(root)
|
||||
if any(prefix in root_name for prefix in exclude_prefix):
|
||||
continue
|
||||
for file in files:
|
||||
|
||||
if file.endswith(('.py', '.js', '.java', '.cpp', '.h', '.c', '.html', '.css')): # add more file types if needed
|
||||
file_path = os.path.join(root, file)
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
code_files.append({
|
||||
"query": f"File: {file_path}\n\nContent:\n{content}",
|
||||
"response": f"This is the content of file {file_path}"
|
||||
})
|
||||
self.add_query(code_files, self.collection_name)
|
||||
|
||||
def query_code(self, query_text: str, n_results: int = 5) -> List[Dict]:
|
||||
"""
|
||||
Query the code memory.
|
||||
|
||||
Args:
|
||||
query_text (str): The query text
|
||||
n_results (int): The number of results to return
|
||||
|
||||
Returns:
|
||||
List[Dict]: The query results list
|
||||
"""
|
||||
results = self.query([query_text], self.collection_name, n_results)
|
||||
return [
|
||||
{
|
||||
"file": doc.split('\n')[0].replace("File: ", ""),
|
||||
"content": '\n'.join(doc.split('\n')[3:]),
|
||||
"metadata": metadata
|
||||
}
|
||||
for doc, metadata in zip(results['documents'][0], results['metadatas'][0])
|
||||
]
|
||||
|
||||
class CodeReranker(Reranker):
|
||||
def __init__(self, model: str) -> None:
|
||||
super().__init__(model)
|
||||
def wrap_query_results(self, query_results: List[Dict]) -> str:
|
||||
wrapped_query_results = ""
|
||||
for result in query_results:
|
||||
wrapped_query_results += f"File: {result['file']}\n"
|
||||
wrapped_query_results += f"Content: {result['content'][:300]}...\n"
|
||||
wrapped_query_results += "---"
|
||||
return wrapped_query_results
|
||||
def wrap_reranked_results(self, reranked_paths: List[str]) -> str:
|
||||
wrapped_reranked_results = "[Referenced code files]:"
|
||||
for path in reranked_paths:
|
||||
wrapped_reranked_results += f"Code path: {path}\n"
|
||||
try:
|
||||
with open(path, 'r', encoding='utf-8') as file:
|
||||
content = file.read()
|
||||
wrapped_reranked_results += f"Code content:\n{content}\n"
|
||||
except Exception as e:
|
||||
wrapped_reranked_results += f"Error reading file: {str(e)}\n"
|
||||
wrapped_reranked_results += "---\n"
|
||||
return wrapped_reranked_results
|
||||
def parse_results(self, reranked_results: str) -> List[str]:
|
||||
lines = reranked_results.strip().split('\n')
|
||||
|
||||
# get the last 5 lines
|
||||
last_lines = lines[-5:]
|
||||
|
||||
# remove the number and dot at the beginning of each line
|
||||
cleaned_lines = [re.sub(r'^\d+\.\s*', '', line.strip()) for line in last_lines]
|
||||
unique_lines = list(dict.fromkeys(cleaned_lines))
|
||||
|
||||
return unique_lines
|
||||
def rerank(self, query_text: str, query_results: List[Dict]) -> List[Dict]:
|
||||
system_prompt = \
|
||||
"""
|
||||
You are a helpful assistant that reranks the given code files (containing the path of files and Overview of the content of files) based on the query.
|
||||
You should rerank the code files based on the query, and the most relevant code files should be ranked on the top.
|
||||
You should select the top 5 code files to answer the query, by giving the file path of the code files.
|
||||
|
||||
Example:
|
||||
[Query]: "The definition of 'BaseAgent'"
|
||||
[Code files]:
|
||||
File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/__init__.py
|
||||
Content: from .ABCAgent import ABCAgent
|
||||
from .BaseAgent import BaseAgent
|
||||
from .ManagerAgent import ManagerAge...
|
||||
---
|
||||
File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/__init__.py
|
||||
Content: from .ABCAgent import ABCAgent
|
||||
from .BaseAgent import BaseAgent
|
||||
from .ManagerAgent import ManagerAge...
|
||||
---
|
||||
File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/__init__.py
|
||||
Content: from .ABCAgent import ABCAgent
|
||||
from .BaseAgent import BaseAgent
|
||||
from .ManagerAgent import ManagerAge...
|
||||
---
|
||||
File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/__init__.py
|
||||
Content: from .ABCAgent import ABCAgent
|
||||
from .BaseAgent import BaseAgent
|
||||
from .ManagerAgent import ManagerAge...
|
||||
---
|
||||
File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agent_prompts/__init__.py
|
||||
Content: from .BasePrompt import BasePromptGen, ManagerPromptGen, PromptGen
|
||||
...
|
||||
---
|
||||
File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agent_prompts/__init__.py
|
||||
Content: from .BasePrompt import BasePromptGen, ManagerPromptGen, PromptGen
|
||||
...
|
||||
---
|
||||
File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agent_prompts/__init__.py
|
||||
Content: from .BasePrompt import BasePromptGen, ManagerPromptGen, PromptGen
|
||||
...
|
||||
---
|
||||
File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agent_prompts/__init__.py
|
||||
Content: from .BasePrompt import BasePromptGen, ManagerPromptGen, PromptGen
|
||||
...
|
||||
---
|
||||
File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/BaseAgent.py
|
||||
Content: from typing import List
|
||||
|
||||
from sa.actions import BaseAction, FinishAct, ThinkAct, PlanAct
|
||||
from sa.age...
|
||||
---
|
||||
File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/BaseAgent.py
|
||||
Content: from typing import List
|
||||
|
||||
from sa.actions import BaseAction, FinishAct, ThinkAct, PlanAct
|
||||
from sa.age...
|
||||
---
|
||||
[Reranked 5 code files]:
|
||||
1. /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/BaseAgent.py
|
||||
2. /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/__init__.py
|
||||
3. /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/ABCAgent.py
|
||||
4. /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/ManagerAgent.py
|
||||
5. /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/AgentLogger.py
|
||||
"""
|
||||
wrapped_query_results = self.wrap_query_results(query_results)
|
||||
user_prompt = \
|
||||
"""
|
||||
[Query]: \n{query_text}
|
||||
[Code files]: \n{query_results}
|
||||
[Reranked 5 code files]:
|
||||
""".format(query_text=query_text, query_results=wrapped_query_results)
|
||||
chat_history = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt}
|
||||
]
|
||||
create_params = {
|
||||
"model": self.model,
|
||||
"messages": chat_history,
|
||||
"stream": False,
|
||||
}
|
||||
response = completion(**create_params)
|
||||
reranked_results = self.parse_results(response.choices[0].message.content)
|
||||
reranked_results = self.wrap_reranked_results(reranked_results)
|
||||
return reranked_results
|
||||
Reference in New Issue
Block a user