Merge pull request #29 from GPTStonks/daedalus/document-openbb-results

Add OpenBB documentation when used
GPTStonks · Dec 26, 2023 · ce63197 · ce63197
2 parents 76899de + 2f171ec
commit ce63197
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 7 deletions.
diff --git a/gptstonks_api/main.py b/gptstonks_api/main.py
@@ -261,12 +261,21 @@ async def run_model_in_background(query: str, use_agent: bool, openbb_pat: str |
             agent_output_str = await agent_executor.arun(query)
 
             try:
-                res = json.loads(agent_output_str)
-                return {
-                    "type": "data",
-                    "result_data": res,
-                    "body": "> Data retrieved using `openbb`. Use with caution.",
-                }
+                if "```json" in agent_output_str:
+                    result_data_str = agent_output_str.split("```json")[1].split("```")[0].strip()
+                    result_data = json.loads(result_data_str)
+                    body_data_str = agent_output_str.split("```json")[0].strip()
+
+                    return {
+                        "type": "data",
+                        "result_data": result_data,
+                        "body": body_data_str,
+                    }
+                else:
+                    return {
+                        "type": "data",
+                        "body": agent_output_str,
+                    }
             except Exception as e:
                 return {
                     "type": "data",

diff --git a/gptstonks_api/utils.py b/gptstonks_api/utils.py
@@ -162,7 +162,27 @@ async def get_openbb_chat_output_executed(
         if "```python" in output_res.response
         else output_res.response
     )
-    return python_repl_utility.run(fix_frequent_code_errors(code_str, openbb_pat))
+    fixed_code_str = fix_frequent_code_errors(code_str, openbb_pat)
+    # run Python and get output
+    repl_output = python_repl_utility.run(fixed_code_str)
+    # get OpenBB's functions called for explicability
+    openbb_funcs_called = []
+    for code_line in code_str.split("\n"):
+        if "obb." in code_line:
+            openbb_funcs_called.append(code_line.split("obb.")[1].strip())
+    openbb_platform_ref_uri = "https://docs.openbb.co/platform/reference/"
+    openbb_funcs_called_str = "".join(
+        [
+            f"- {x} [[documentation]({openbb_platform_ref_uri}{x.split('(')[0].replace('.', '/')})]\n"
+            for x in openbb_funcs_called
+        ]
+    )
+
+    return (
+        "> Context retrieved using OpenBB. "
+        f"OpenBB's functions called:\n{openbb_funcs_called_str.strip()}\n\n"
+        f"```json\n{repl_output.strip()}\n```"
+    )
 
 
 def run_qa_over_tool_output(tool_input: str | dict, llm: BaseLLM, tool: BaseTool) -> str: