Merge pull request #75 from GPTStonks/enh/add-flexibility-to-mitheith…

…el-search-tool

Add more parameters to Mitheithel Search tool

libs/gptstonks-multiagents/gptstonks/multiagents/tools/mitheithel_search.py

@@ -1,10 +1,12 @@
+import json
 import os
+import urllib.parse
 import warnings
 from functools import partial
 
-import httpx
 from langchain.tools import StructuredTool
 from pydantic import BaseModel, Field
+from websockets.asyncio.client import connect
 
 try:
     MITHEITHEL_API_KEY = os.environ["MITHEITHEL_API_KEY"]
@@ -22,18 +24,41 @@ class MitheithelSearchInput(BaseModel):
         query: str = Field(description="query for the queries to search with Mitheithel.")
 
     @classmethod
-    async def search(cls, query: str, timeout: float = 180) -> str:
-        """Searches videos on Youtube related to a query."""
-
-        async with httpx.AsyncClient(timeout=timeout) as client:
-            mith_res = await client.get(
-                url=MITHEITHEL_SEARCH_URI,
-                params={
-                    "query": query,
-                },
-                headers={"Authorization": f"Bearer {MITHEITHEL_API_KEY}"},
-            )
-            mith_res_data: dict = mith_res.json()
+    async def search(
+        cls,
+        query: str,
+        timeout: float = 180,
+        use_quality: bool = False,
+        timelimit: str = "w",
+        return_json: bool = False,
+    ) -> str:
+        """Searches videos on Youtube related to a query.
+
+        Args:
+            query (`str`): what to search using Mitheithel agents. It works better with full requests in natural language, instead of keywords.
+            timeout (`float`): maximum time to wait for a response, in seconds. Speed takes 15-30s, and Quality 120-150s.
+            use_quality (`bool`): whether to use quality or speed mode when searching.
+            timelimit (`str`): how far into the past our search engine can look for data. One of d, w, m, y. Defaults to one week old.
+            return_json (`bool`): whether to return the complete JSON or only the content.
+
+        Returns:
+            `str`: the complete JSON from Mitheithel API or only its content, depending on `return_json`.
+        """
+
+        params: dict = {
+            "token": MITHEITHEL_API_KEY,
+            "use_quality": use_quality,
+            "timelimit": timelimit,
+        }
+        websocket_uri: str = (
+            f"{MITHEITHEL_SEARCH_URI.replace('https', 'wss')}?{urllib.parse.urlencode(params)}"
+        )
+        async with connect(websocket_uri) as websocket:
+            await websocket.send(json.dumps({"query": query}))
+            mith_res: str = await websocket.recv()
+        if return_json:
+            return mith_res
+        mith_res_data: dict = json.loads(mith_res)
         references: str = "\n".join(f"- {x}" for x in mith_res_data["references"])
         return f'Content: {mith_res_data["body"]}\nReferences:\n{references}\n'
 
@@ -44,10 +69,33 @@ def create(
         description: str = "Useful to search anything on the Internet",
         return_direct: bool = False,
         timeout: float = 180,
+        use_quality: bool = False,
+        timelimit: str = "w",
+        return_json: bool = False,
     ) -> StructuredTool:
+        """Creates a LangChain tool to interact with Mitheithel API.
+
+        Args:
+            name (`str`): tool name for LangChain.
+            name (`str`): the purpose of the tool, according to LangChain's documentation.
+            return_direct (`bool`): Whether to return the result directly or as a callback.
+            timeout (`float`): maximum time to wait for a response, in seconds. Speed takes 15-30s, and Quality 120-150s.
+            use_quality (`bool`): whether to use quality or speed mode when searching.
+            timelimit (`str`): how far into the past our search engine can look for data. One of d, w, m, y. Defaults to one week old.
+            return_json (`bool`): whether to return the complete JSON or only the content.
+
+        Returns:
+            `StructuredTool`: the tool."""
+
         return cls.from_function(
             func=None,
-            coroutine=partial(cls.search, timeout=timeout),
+            coroutine=partial(
+                cls.search,
+                timeout=timeout,
+                use_quality=use_quality,
+                timelimit=timelimit,
+                return_json=return_json,
+            ),
             name=name,
             description=description,
             args_schema=cls.MitheithelSearchInput,

libs/gptstonks-multiagents/pyproject.toml

@@ -13,7 +13,7 @@ classifiers = [
   "Operating System :: OS Independent",
 ]
 dependencies = [
-    "gptstonks-wrappers>=0.0.1.post3",
+    "gptstonks-wrappers>=0.0.3.post1",
     "langgraph>=0.2.20,<0.3",
     "langchain>=0.3,<0.4",
     "asyncio>=3.4.3",
@@ -22,6 +22,7 @@ dependencies = [
     "wikipedia>=1.4.0",
     "llama-index-llms-openai>=0.1.16",
     "langchain-community<0.4,>=0.3",
+    "websockets>=14.1",
 ]
 license = {text = "MIT"}
 

libs/gptstonks-multiagents/tests/tools/test_mitheithel_search.py

@@ -1,19 +1,22 @@
-from unittest.mock import AsyncMock, patch
+import json
+from unittest.mock import patch
 
-import httpx
 import pytest
 from langchain.tools import StructuredTool
+from websockets.asyncio.client import connect
 
-from gptstonks.multiagents.tools import MitheithelSearchTool
 
+@pytest.mark.asyncio
+async def test_mith_search():
+    with (patch.object(connect, "create_connection") as mock_connect,):
+        from gptstonks.multiagents.tools import MitheithelSearchTool
 
-@patch.object(
-    httpx._client.AsyncClient,
-    "get",
-    AsyncMock(
-        return_value=httpx.Response(
-            200,
-            json={
+        mith_search_tool = MitheithelSearchTool.create()
+        assert isinstance(mith_search_tool, StructuredTool)
+        with patch.object(
+            json,
+            "loads",
+            return_value={
                 "type": "string",
                 "body": "string",
                 "follow_up_questions": ["string"],
@@ -23,13 +26,7 @@
                 "subqueries_answered": ["string"],
                 "subqueries_responses": ["string"],
             },
-        )
-    ),
-)
-@pytest.mark.asyncio
-async def test_mith_search():
-    mith_search_tool = MitheithelSearchTool.create()
-    assert isinstance(mith_search_tool, StructuredTool)
-    res = await mith_search_tool.ainvoke("gptstonks")
-    assert isinstance(res, str)
-    assert "References:" in res
+        ):
+            res = await mith_search_tool.ainvoke("gptstonks")
+            assert isinstance(res, str)
+            assert "References:" in res