Merge pull request #189 from KTC-Security-Circle/feature/retrun-only-…

…source-id

source_idのみ返却に変更

src/sc_system_ai/agents/classify_agent.py

@@ -64,7 +64,7 @@ def invoke(self, message: str) -> AgentResponse:
         resp.document_id = self._doc_id_checker()
         return resp
 
-    def _doc_id_checker(self) -> list[str | int] | None:
+    def _doc_id_checker(self) -> list[int] | None:
         """
         ソースIDが存在するか確認する
         """

src/sc_system_ai/agents/search_school_data_agent.py

@@ -34,16 +34,14 @@ def __init__(
         )
         self.assistant_info = search_school_data_agent_info
 
-    def _add_search_result(self, message: str) -> list[str | int]:
+    def _add_search_result(self, message: str) -> list[int]:
         word = genarate_search_word(message)
         search = search_school_database_cosmos(word)
         ids = []
         for doc in search:
             self.assistant_info += f"### {doc.metadata['title']}\n" + doc.page_content + "\n"
             if "source_id" in doc.metadata:
                 ids.append(doc.metadata["source_id"])
-            else:
-                ids.append(doc.metadata["id"])
         super().set_assistant_info(self.assistant_info)
         return ids
 

src/sc_system_ai/agents/tools/calling_search_school_data_agent.py

@@ -11,7 +11,7 @@
 
 class CallingSearchSchoolDataAgent(CallingAgent):
     # tool側でidを保持する
-    source_id: set[str | int] = set()
+    source_id: set[int] = set()
 
     def __init__(self) -> None:
         super().__init__()

src/sc_system_ai/main.py

@@ -73,7 +73,7 @@
 class Response(TypedDict):
     output: str | None
     error: str | None
-    document_id: list[str | int] | None
+    document_id: list[int] | None
 
 class StreamResponse(TypedDict):
     output: str | None

src/sc_system_ai/template/agent.py

@@ -102,7 +102,7 @@ class AgentResponse(BaseAgentResponse):
     """Agentのレスポンスの型"""
     chat_history: list[HumanMessage | AIMessage] | None = None
     messages: str | None = None
-    document_id: list[str | int] | None = None
+    document_id: list[int] | None = None
 
 
 class StreamingAgentResponse(BaseAgentResponse):