Skip to content

Commit

Permalink
making the prompt generation a function inside query_filter.py
Browse files Browse the repository at this point in the history
  • Loading branch information
Jgmedina95 committed Mar 18, 2024
1 parent 43aeb62 commit c51547e
Show file tree
Hide file tree
Showing 2 changed files with 88 additions and 87 deletions.
89 changes: 2 additions & 87 deletions mdagent/mainagent/agent.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
import json

from dotenv import load_dotenv
from langchain.agents import AgentExecutor, OpenAIFunctionsAgent
from langchain.agents.structured_chat.base import StructuredChatAgent
Expand All @@ -10,8 +8,7 @@
from mdagent.utils import PathRegistry, _make_llm

from ..tools import get_tools, make_all_tools
from .prompt import modular_analysis_prompt, openaifxn_prompt, structured_prompt
from .query_filter import Parameters, Task_type, create_filtered_query
from .query_filter import make_prompt

load_dotenv()

Expand Down Expand Up @@ -122,88 +119,6 @@ def _initialize_tools_and_agent(self, user_input=None):
)

def run(self, user_input, callbacks=None):
if self.agent_type == "Structured":
tries = 1

while tries <= 3:
try:
structured_query = create_filtered_query(
user_input, model="gpt-3.5-turbo"
)
structured_query = json.loads(structured_query)
parameters = Parameters.parse_parameters_string(
structured_query["Parameters"]
)
_parameters = ""
for key, value in parameters.items():
if value == "None":
continue
else:
_parameters += f"{key}: {value}, "
_plan = ""
if structured_query["UserProposedPlan"] == "[]":
_plan += "None"
else:
if type(structured_query["UserProposedPlan"]) == str:
for plan in structured_query["UserProposedPlan"].split(","):
_plan += f"{plan},"
elif type(structured_query["UserProposedPlan"]) == list:
for plan in structured_query["UserProposedPlan"]:
_plan += f"{plan},"
_proteins = ""
if structured_query["ProteinS"] == "['None']":
_proteins += "None"
elif structured_query["ProteinS"] == "[]":
_proteins += "None"
else:
for protein in eval(structured_query["ProteinS"]):
_proteins += f"{protein}, "
_subtasks = ""
if structured_query["Subtask_types"] == "['None']":
_subtasks += "None"
elif structured_query["Subtask_types"] == "[]":
_subtasks += "None"
elif structured_query["Subtask_types"] == ["None"]:
_subtasks += "None"
else:
if type(structured_query["Subtask_types"]) == str:
for subtask in Task_type.parse_task_type_string(
structured_query["Subtask_types"]
):
_subtasks += f"{subtask}, "
elif type(structured_query["Subtask_types"]) == list:
for subtask in structured_query["Subtask_types"]:
_str = Task_type.parse_task_type_string(subtask)
_subtasks += f"{_str}, "
prompt = modular_analysis_prompt.format(
Main_Task=structured_query["Main_Task"],
Subtask_types=_subtasks,
Proteins=_proteins,
Parameters=_parameters,
UserProposedPlan=_plan,
)
break
except ValueError as e:
print(f"Failed to structure query, attempt {tries}/3. Retrying...")
print(e, e.args)
tries += 1
continue
except Exception as e:
print(f"Failed to structure query, attempt {tries}/3. Retrying...")
print(e, e.args)
tries += 1
continue

if tries > 3:
print(
"Failed to structure query after 3 attempts."
"Input will be used as is."
)
self.prompt = structured_prompt.format(input=user_input)
else:
self.prompt = prompt
elif self.agent_type == "OpenAIFunctionsAgent":
self.prompt = openaifxn_prompt.format(input=user_input)

self.prompt = make_prompt(user_input, self.agent_type, model="gpt-3.5-turbo")
self.agent = self._initialize_tools_and_agent(user_input)
return self.agent.run(self.prompt, callbacks=callbacks)
86 changes: 86 additions & 0 deletions mdagent/mainagent/query_filter.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import json
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional
Expand All @@ -6,6 +7,8 @@
from outlines import generate, models
from pydantic import BaseModel

from .prompt import modular_analysis_prompt, openaifxn_prompt, structured_prompt

################################################################

"""
Expand Down Expand Up @@ -252,3 +255,86 @@ def create_filtered_query(raw_query, model="gpt-3.5-turbo", examples=examples):
filter_model = models.openai(model)
generator = generate.text(filter_model)
return generator(query_filter(raw_query, examples=examples))


def make_prompt(user_input, agent_type, model="gpt-3.5-turbo"):
if agent_type == "Structured":
tries = 1

while tries <= 3:
try:
structured_query = create_filtered_query(user_input, model=model)
structured_query = json.loads(structured_query)
parameters = Parameters.parse_parameters_string(
structured_query["Parameters"]
)
_parameters = ""
for key, value in parameters.items():
if value == "None":
continue
else:
_parameters += f"{key}: {value}, "
_plan = ""
if structured_query["UserProposedPlan"] == "[]":
_plan += "None"
else:
if type(structured_query["UserProposedPlan"]) == str:
for plan in structured_query["UserProposedPlan"].split(","):
_plan += f"{plan},"
elif type(structured_query["UserProposedPlan"]) == list:
for plan in structured_query["UserProposedPlan"]:
_plan += f"{plan},"
_proteins = ""
if structured_query["ProteinS"] == "['None']":
_proteins += "None"
elif structured_query["ProteinS"] == "[]":
_proteins += "None"
else:
for protein in eval(structured_query["ProteinS"]):
_proteins += f"{protein}, "
_subtasks = ""
if structured_query["Subtask_types"] == "['None']":
_subtasks += "None"
elif structured_query["Subtask_types"] == "[]":
_subtasks += "None"
elif structured_query["Subtask_types"] == ["None"]:
_subtasks += "None"
else:
if type(structured_query["Subtask_types"]) == str:
for subtask in Task_type.parse_task_type_string(
structured_query["Subtask_types"]
):
_subtasks += f"{subtask}, "
elif type(structured_query["Subtask_types"]) == list:
for subtask in structured_query["Subtask_types"]:
_str = Task_type.parse_task_type_string(subtask)
_subtasks += f"{_str}, "
prompt = modular_analysis_prompt.format(
Main_Task=structured_query["Main_Task"],
Subtask_types=_subtasks,
Proteins=_proteins,
Parameters=_parameters,
UserProposedPlan=_plan,
)
break
except ValueError as e:
print(f"Failed to structure query, attempt {tries}/3. Retrying...")
print(e, e.args)
tries += 1
continue
except Exception as e:
print(f"Failed to structure query, attempt {tries}/3. Retrying...")
print(e, e.args)
tries += 1
continue

if tries > 3:
print(
"Failed to structure query after 3 attempts."
"Input will be used as is."
)
return structured_prompt.format(input=user_input)
else:
return prompt
elif agent_type == "OpenAIFunctionsAgent":
return openaifxn_prompt.format(input=user_input)

0 comments on commit c51547e

Please sign in to comment.