Chatクラスからの呼び出し処理を作成

src/sc_system_ai/main.py

@@ -60,7 +60,7 @@
 import logging
 from collections.abc import Iterator
 from importlib import import_module
-from typing import Literal, TypedDict
+from typing import Literal, TypedDict, cast
 
 from sc_system_ai.template.agent import Agent
 from sc_system_ai.template.ai_settings import llm
@@ -144,7 +144,7 @@ def invoke(
         self,
         message: str,
         command: AGENT = "classify"
-    ) -> Iterator[str]:
+    ) -> Iterator[Response]:
         """エージェントを呼び出し、チャットを行う関数
 
         Args:
@@ -167,20 +167,15 @@ def invoke(
         if self.is_streaming:
             for resp in self.agent.invoke(message):
                 if type(resp) is str:
-                    yield resp
+                    yield self._create_response({"output": resp})
         else:
             resp = next(self.agent.invoke(message))
-
-            if type(resp) is dict:
-                if "error" in resp:
-                    yield resp["error"]
-                else:
-                    yield resp["output"]
+            yield self._create_response(cast(dict, resp))
 
     def _call_agent(self, command: AGENT) -> None:
         try:
             module_name = f"sc_system_ai.agents.{command}_agent"
-            class_name = f"{command.capitalize()}Agent"
+            class_name = "".join([cn.capitalize() for cn in command.split("_")]) + "Agent"
             module = import_module(module_name)
             agent_class = getattr(module, class_name)
 
@@ -194,6 +189,13 @@ def _call_agent(self, command: AGENT) -> None:
             logger.error(f"エージェントが見つかりません: {command}")
             raise ValueError(f"エージェントが見つかりません: {command}") from None
 
+    def _create_response(self, resp: dict) -> Response:
+        return {
+            "output": resp.get("output"),
+            "error": resp.get("error"),
+            "document_id": resp.get("document_id")
+        }
+
 
 
 def static_chat() -> None:
@@ -257,15 +259,15 @@ def streaming_chat() -> None:
         ],
         is_streaming=False,
     )
-    message = "京都テックについて教えて"
+    message = "私の名前と専攻は何ですか？"
 
     # try:
     #     resp = chat.agent.get_response()
     # except Exception:
     #     pass
 
     # # 通常呼び出し
-    resp = next(chat.invoke(message=message, command="search_school_data"))
+    resp = next(chat.invoke(message=message, command="dummy"))
     print(resp)
 
     # ストリーミング呼び出し