Switch from Azure OpenAI to Bedrock

nulib · Dec 12, 2024 · a075cfa · a075cfa
1 parent a8effff
commit a075cfa
Show file tree

Hide file tree

Showing 9 changed files with 21 additions and 391 deletions.
diff --git a/chat/src/agent/agent_handler.py b/chat/src/agent/agent_handler.py
@@ -52,7 +52,8 @@ def on_tool_end(self, output: ToolMessage, **kwargs: Dict[str, Any]):
                 pass
             case "search":
                 try:
-                    docs: List[Dict[str, Any]] = [doc.metadata for doc in output.artifact]
+                    result_fields = ("id", "title", "visibility", "work_type", "thumbnail")
+                    docs: List[Dict[str, Any]] = [{k: doc.metadata.get(k) for k in result_fields} for doc in output.artifact]
                     self.socket.send({"type": "search_result", "ref": self.ref, "message": docs})
                 except json.decoder.JSONDecodeError as e:
                     print(f"Invalid json ({e}) returned from {output.name} tool: {output.content}")

diff --git a/chat/src/agent/dynamodb_cleaner.py b/chat/src/agent/dynamodb_cleaner.py
diff --git a/chat/src/agent/dynamodb_saver.py b/chat/src/agent/dynamodb_saver.py