46 lines
2.3 KiB
Python
46 lines
2.3 KiB
Python
|
|
from coagent.connector.schema import Message
|
|
from .prompt_manager import PromptManager
|
|
|
|
|
|
class Code2DocPM(PromptManager):
|
|
def handle_code_snippet(self, **kwargs) -> str:
|
|
if 'previous_agent_message' not in kwargs:
|
|
return ""
|
|
previous_agent_message: Message = kwargs['previous_agent_message']
|
|
code_snippet = previous_agent_message.customed_kargs.get("Code Snippet", "")
|
|
current_vertex = previous_agent_message.customed_kargs.get("Current_Vertex", "")
|
|
instruction = "A segment of code that contains the function or method to be documented.\n"
|
|
return instruction + "\n" + f"name: {current_vertex}\n{code_snippet}"
|
|
|
|
def handle_specific_objective(self, **kwargs) -> str:
|
|
if 'previous_agent_message' not in kwargs:
|
|
return ""
|
|
previous_agent_message: Message = kwargs['previous_agent_message']
|
|
specific_objective = previous_agent_message.parsed_output.get("Code Path")
|
|
|
|
instruction = "Provide the code path of the function or method you wish to document.\n"
|
|
s = instruction + f"\n{specific_objective}"
|
|
return s
|
|
|
|
|
|
class CodeRetrievalPM(PromptManager):
|
|
def handle_code_snippet(self, **kwargs) -> str:
|
|
if 'previous_agent_message' not in kwargs:
|
|
return ""
|
|
previous_agent_message: Message = kwargs['previous_agent_message']
|
|
code_snippet = previous_agent_message.customed_kargs.get("Code Snippet", "")
|
|
current_vertex = previous_agent_message.customed_kargs.get("Current_Vertex", "")
|
|
instruction = "the initial Code or objective that the user wanted to achieve"
|
|
return instruction + "\n" + f"name: {current_vertex}\n{code_snippet}"
|
|
|
|
def handle_retrieval_codes(self, **kwargs) -> str:
|
|
if 'previous_agent_message' not in kwargs:
|
|
return ""
|
|
previous_agent_message: Message = kwargs['previous_agent_message']
|
|
Retrieval_Codes = previous_agent_message.customed_kargs["Retrieval_Codes"]
|
|
Relative_vertex = previous_agent_message.customed_kargs["Relative_vertex"]
|
|
instruction = "the initial Code or objective that the user wanted to achieve"
|
|
s = instruction + "\n" + "\n".join([f"name: {vertext}\n{code}" for vertext, code in zip(Relative_vertex, Retrieval_Codes)])
|
|
return s
|