Merge pull request #389 from kreneskyp/knowledge_fix

Fixing Knowledge agent:
kreneskyp · Jan 15, 2024 · 782ccfb · 782ccfb
2 parents 331e188 + 28430cb
commit 782ccfb
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 22 deletions.
diff --git a/ix/chains/fixtures/agent/knowledge.json b/ix/chains/fixtures/agent/knowledge.json
@@ -76,8 +76,8 @@
       "description": "",
       "root": true,
       "position": {
-        "x": 150.0,
-        "y": 700.0
+        "x": -140.0,
+        "y": 680.0
       },
       "chain": "1a427ad9-e764-4f99-be8e-bf55bea1d9e8"
     }
@@ -119,7 +119,8 @@
         "ix.chains.components.json.JSONPath"
       ],
       "config": {
-        "path": "output"
+        "path": "output",
+        "return_list": true
       },
       "name": "",
       "description": "",
@@ -222,6 +223,28 @@
       "chain": "1a427ad9-e764-4f99-be8e-bf55bea1d9e8"
     }
   },
+  {
+    "model": "chains.chainnode",
+    "pk": "eef41348-4017-4bcd-bcbb-f81e4b277eef",
+    "fields": {
+      "class_path": "ix.chains.components.json.JSONPath",
+      "node_type": [
+        "ix.chains.components.json.JSONPath"
+      ],
+      "config": {
+        "path": "$.user_input",
+        "return_list": false
+      },
+      "name": "",
+      "description": "",
+      "root": false,
+      "position": {
+        "x": 210.0,
+        "y": 760.0
+      },
+      "chain": "1a427ad9-e764-4f99-be8e-bf55bea1d9e8"
+    }
+  },
   {
     "model": "chains.chainnode",
     "pk": "fbd0e538-59c3-4b1d-bd5f-95e3d536b08d",
@@ -237,7 +260,7 @@
       "description": "",
       "root": false,
       "position": {
-        "x": 470.0,
+        "x": 210.0,
         "y": 640.0
       },
       "chain": "1a427ad9-e764-4f99-be8e-bf55bea1d9e8"
@@ -280,14 +303,27 @@
       "chain": "1a427ad9-e764-4f99-be8e-bf55bea1d9e8"
     }
   },
+  {
+    "model": "chains.chainedge",
+    "pk": "01cde37e-9e9e-4f78-8712-10b41f660ac8",
+    "fields": {
+      "source": "eef41348-4017-4bcd-bcbb-f81e4b277eef",
+      "target": "0d49c55f-896e-4dc8-9bf6-36e8057adfeb",
+      "source_key": "out",
+      "target_key": "in",
+      "chain": "1a427ad9-e764-4f99-be8e-bf55bea1d9e8",
+      "input_map": null,
+      "relation": "LINK"
+    }
+  },
   {
     "model": "chains.chainedge",
     "pk": "06c25dac-3972-4cbe-a798-748095ccda89",
     "fields": {
       "source": "a3274983-debb-4e34-b1eb-ff758b22a547",
       "target": "e4bca73a-e160-4182-9984-a9594cf2149a",
       "source_key": "embeddings",
-      "target_key": "embedding",
+      "target_key": "embedding_function",
       "chain": "1a427ad9-e764-4f99-be8e-bf55bea1d9e8",
       "input_map": null,
       "relation": "PROP"
@@ -345,19 +381,6 @@
       "relation": "LINK"
     }
   },
-  {
-    "model": "chains.chainedge",
-    "pk": "9b3d1155-17ba-4440-a104-8fd6281eb58a",
-    "fields": {
-      "source": "43e1d43e-4117-41c8-bbee-26efbff67ea3",
-      "target": "0d49c55f-896e-4dc8-9bf6-36e8057adfeb",
-      "source_key": "chat_input",
-      "target_key": "in",
-      "chain": "1a427ad9-e764-4f99-be8e-bf55bea1d9e8",
-      "input_map": null,
-      "relation": "LINK"
-    }
-  },
   {
     "model": "chains.chainedge",
     "pk": "bd960a30-7632-456e-a050-626f3abdfd45",
@@ -399,11 +422,11 @@
   },
   {
     "model": "chains.chainedge",
-    "pk": "caa53c20-976d-4547-b2cc-c0a81dff34c8",
+    "pk": "e7a43957-99b1-4b4b-bdb1-3f7680c1b2ad",
     "fields": {
-      "source": "fbd0e538-59c3-4b1d-bd5f-95e3d536b08d",
-      "target": "4e15524b-1cb1-4960-8fea-e29f765ba751",
-      "source_key": "out",
+      "source": "43e1d43e-4117-41c8-bbee-26efbff67ea3",
+      "target": "eef41348-4017-4bcd-bcbb-f81e4b277eef",
+      "source_key": "chat_input",
       "target_key": "in",
       "chain": "1a427ad9-e764-4f99-be8e-bf55bea1d9e8",
       "input_map": null,

diff --git a/ix/chains/tests/agents/__init__.py b/ix/chains/tests/agents/__init__.py
diff --git a/ix/chains/tests/agents/test_knowledge.py b/ix/chains/tests/agents/test_knowledge.py
@@ -0,0 +1,26 @@
+import pytest
+from langchain_core.messages import AIMessage
+
+from ix.chains.loaders.core import ainit_chain_flow
+from ix.chains.models import Chain
+from ix.conftest import aload_fixture
+
+
+@pytest.mark.openai_api
+@pytest.mark.django_db
+class TestKnowledge:
+    """
+    Tests for Knowledge agent.
+    """
+
+    async def test_knowledge_flow(self, anode_types, aix_context):
+        """test loading the knowledge agent"""
+
+        await aload_fixture("agent/knowledge")
+        chain = await Chain.objects.aget(agent__alias="knowledge")
+
+        # init flow
+        runnable = await ainit_chain_flow(chain, context=aix_context)
+
+        output = await runnable.ainvoke(input={"user_input": "test"})
+        assert isinstance(output, AIMessage)