[AUR-431, AUR-435] Add Agent Interface and ReWOO Agent implementation (#31)

* add base Tool * minor update test_tool * update test dependency * update test dependency * Fix namespace conflict * update test * add base Agent Interface, add ReWoo Agent * minor update * update test * fix typo * remove unneeded print * update rewoo agent --------- Co-authored-by: trducng <trungduc1992@gmail.com>
2023-10-01 11:53:08 +07:00
parent f9fc02a32a
commit 91048770fa
11 changed files with 919 additions and 0 deletions
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -0,0 +1,68 @@
+from unittest.mock import patch
+
+from kotaemon.llms.chats.openai import AzureChatOpenAI
+from kotaemon.pipelines.agents.rewoo import RewooAgent
+from kotaemon.pipelines.tools import GoogleSearchTool, WikipediaTool
+
+FINAL_RESPONSE_TEXT = "Hello Cinnamon AI!"
+_openai_chat_completion_responses = [
+    {
+        "id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
+        "object": "chat.completion",
+        "created": 1692338378,
+        "model": "gpt-35-turbo",
+        "choices": [
+            {
+                "index": 0,
+                "finish_reason": "stop",
+                "message": {
+                    "role": "assistant",
+                    "content": "#Plan1: Search for Cinnamon AI company on Google\n"
+                    "#E1: google_search[Cinnamon AI company]\n"
+                    "#Plan2: Search for Cinnamon on Wikipedia\n"
+                    "#E2: wikipedia[Cinnamon]",
+                },
+            }
+        ],
+        "usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
+    },
+    {
+        "id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
+        "object": "chat.completion",
+        "created": 1692338378,
+        "model": "gpt-35-turbo",
+        "choices": [
+            {
+                "index": 0,
+                "finish_reason": "stop",
+                "message": {
+                    "role": "assistant",
+                    "content": FINAL_RESPONSE_TEXT,
+                },
+            }
+        ],
+        "usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
+    },
+]
+
+
+@patch(
+    "openai.api_resources.chat_completion.ChatCompletion.create",
+    side_effect=_openai_chat_completion_responses,
+)
+def test_rewoo_agent(openai_completion):
+    llm = AzureChatOpenAI(
+        openai_api_base="https://dummy.openai.azure.com/",
+        openai_api_key="dummy",
+        openai_api_version="2023-03-15-preview",
+        deployment_name="dummy-q2",
+        temperature=0,
+    )
+
+    plugins = [GoogleSearchTool(), WikipediaTool()]
+
+    agent = RewooAgent(llm=llm, plugins=plugins)
+
+    response = agent("Tell me about Cinnamon AI company")
+    openai_completion.assert_called()
+    assert response.output == FINAL_RESPONSE_TEXT