Add Citation pipeline (#78)

* add rerankers in retrieving pipeline

* update example MVP pipeline

* add citation pipeline and function call interface

* change return type of QA and AgentPipeline to Document

tests/test_agent.py

@@ -163,7 +163,7 @@ def test_rewoo_agent(openai_completion, llm, mock_google_search):
 
     response = agent("Tell me about Cinnamon AI company")
     openai_completion.assert_called()
-    assert response.output == FINAL_RESPONSE_TEXT
+    assert response.text == FINAL_RESPONSE_TEXT
 
 
 @patch(
@@ -180,7 +180,7 @@ def test_react_agent(openai_completion, llm, mock_google_search):
 
     response = agent("Tell me about Cinnamon AI company")
     openai_completion.assert_called()
-    assert response.output == FINAL_RESPONSE_TEXT
+    assert response.text == FINAL_RESPONSE_TEXT
 
 
 @patch(
@@ -224,4 +224,4 @@ def test_react_agent_with_langchain_tools(openai_completion, llm):
 
     response = agent("Tell me about Cinnamon AI company")
     openai_completion.assert_called()
-    assert response.output == FINAL_RESPONSE_TEXT
+    assert response.text == FINAL_RESPONSE_TEXT

tests/test_citation.py

@@ -0,0 +1,61 @@
+# flake8: noqa
+from unittest.mock import patch
+
+import pytest
+from openai.types.chat.chat_completion import ChatCompletion
+
+from kotaemon.llms.chats.openai import AzureChatOpenAI
+from kotaemon.pipelines.citation import CitationPipeline
+
+function_output = '{\n  "question": "What is the provided _example_ benefits?",\n  "answer": [\n    {\n      "fact": "特約死亡保険金: 被保険者がこの特約の保険期間中に死亡したときに支払います。",\n      "substring_quote": ["特約死亡保険金"]\n    },\n    {\n      "fact": "特約特定疾病保険金: 被保険者がこの特約の保険期間中に特定の疾病（悪性新生物（がん）、急性心筋梗塞または脳卒中）により所定の状態に該当したときに支払います。",\n      "substring_quote": ["特約特定疾病保険金"]\n    },\n    {\n      "fact": "特約障害保険金: 被保険者がこの特約の保険期間中に傷害もしくは疾病により所定の身体障害の状態に該当したとき、または不慮の事故により所定の身体障害の状態に該当したときに支払います。",\n      "substring_quote": ["特約障害保険金"]\n    },\n    {\n      "fact": "特約介護保険金: 被保険者がこの特約の保険期間中に傷害または疾病により所定の要介護状態に該当したときに支払います。",\n      "substring_quote": ["特約介護保険金"]\n    }\n  ]\n}'
+
+_openai_chat_completion_response = [
+    ChatCompletion.parse_obj(
+        {
+            "id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
+            "object": "chat.completion",
+            "created": 1692338378,
+            "model": "gpt-35-turbo",
+            "system_fingerprint": None,
+            "choices": [
+                {
+                    "index": 0,
+                    "finish_reason": "function_call",
+                    "message": {
+                        "role": "assistant",
+                        "content": None,
+                        "function_call": {
+                            "arguments": function_output,
+                            "name": "QuestionAnswer",
+                        },
+                        "tool_calls": None,
+                    },
+                }
+            ],
+            "usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
+        }
+    )
+]
+
+
+@pytest.fixture
+def llm():
+    return AzureChatOpenAI(
+        azure_endpoint="https://dummy.openai.azure.com/",
+        openai_api_key="dummy",
+        openai_api_version="2023-03-15-preview",
+        temperature=0,
+    )
+
+
+@patch(
+    "openai.resources.chat.completions.Completions.create",
+    side_effect=_openai_chat_completion_response,
+)
+def test_citation(openai_completion, llm):
+    question = "test query"
+    context = "document context"
+
+    citation = CitationPipeline(llm=llm)
+    result = citation(context, question)
+    assert len(result.answer) == 4