Merge pull request #28 from GPTStonks/daedalus/improve-qa-with-comple…

…te-query

Add the complete query to the QA tool

gptstonks_api/main.py

@@ -239,6 +239,9 @@ async def run_model_in_background(query: str, use_agent: bool, openbb_pat: str |
         if use_agent:
             # update openbb tool with PAT (None if not provided)
             app.tools[-1].coroutine = partial(app.tools[-1].coroutine, openbb_pat=openbb_pat)
+            # update QA tools to use original query to respond from the context
+            app.tools[0].coroutine = partial(app.tools[0].coroutine, original_query=query)
+            app.tools[1].coroutine = partial(app.tools[1].coroutine, original_query=query)
             agent_executor = initialize_agent(
                 tools=app.tools,
                 llm=app.llm,

gptstonks_api/utils.py

@@ -176,13 +176,22 @@ def run_qa_over_tool_output(tool_input: str | dict, llm: BaseLLM, tool: BaseTool
     return f"> Context retrieved using {tool.name}.\n\n" f"{answer}"
 
 
-async def arun_qa_over_tool_output(tool_input: str | dict, llm: BaseLLM, tool: BaseTool) -> str:
+async def arun_qa_over_tool_output(
+    tool_input: str | dict, llm: BaseLLM, tool: BaseTool, original_query: Optional[str] = None
+) -> str:
     tool_output: str = await tool.arun(tool_input)
-    model_prompt: str = PromptTemplate(
+    model_prompt = PromptTemplate(
         input_variables=["context_str", "query_str"],
         template=os.getenv("CUSTOM_GPTSTONKS_QA", DEFAULT_TEXT_QA_PROMPT_TMPL),
-    ).format(query_str=tool_input, context_str=tool_output)
-    answer: str = await llm.apredict(model_prompt)
+    )
+    if original_query is not None:
+        answer: str = await llm.apredict(
+            model_prompt.format(query_str=original_query, context_str=tool_output)
+        )
+    else:
+        answer: str = await llm.apredict(
+            model_prompt.format(query_str=tool_input, context_str=tool_output)
+        )
 
     return f"> Context retrieved using {tool.name}.\n\n" f"{answer}"