Update orchestrator/dispatcher.py

orchestrator/dispatcher.py

@@ -1,53 +1,52 @@
 # File: orchestrator/dispatcher.py
-import uuid
 import yaml
 from orchestrator.client import MCPClient
 
-
 class Dispatcher:
     """
-    Coordinates calls to various MCP servers for searching, code execution, and graph retrieval.
+    Coordinates calls to MCP servers for search, notebook execution, and graph retrieval.
     """
-    def __init__(self, config_path="config.yaml"):
-        cfg = yaml.safe_load(open(config_path))
-        servers = cfg.get("mcp_servers", {})
-        self.web = MCPClient(servers.get("web_search"))
-        self.pubmed = MCPClient(servers.get("pubmed"))
-        self.chroma = MCPClient(servers.get("chroma"))
-        self.runner = MCPClient(servers.get("python_run"))
+    def __init__(self, config_path='config.yaml'):
+        with open(config_path, 'r') as f:
+            cfg = yaml.safe_load(f)
+        servers = cfg.get('mcp_servers', {})
+        self.web = MCPClient(servers.get('web_search', ''))
+        self.pubmed = MCPClient(servers.get('pubmed', ''))
+        self.chroma = MCPClient(servers.get('chroma', ''))
+        self.runner = MCPClient(servers.get('python_run', ''))
 
     def search_papers(self, query: str, limit: int = 5):
         """
-        Fan out search requests to web and PubMed MCP servers, aggregate and limit.
+        Aggregate results from web and PubMed MCP endpoints.
         """
-        results = []
+        papers = []
         try:
-            results += self.web.call("web_search.search", {"q": query})
+            papers += self.web.call('web_search.search', {'q': query}) or []
         except Exception:
             pass
         try:
-            results += self.pubmed.call("metatool.query", {"source": "PubMed", "q": query})
+            papers += self.pubmed.call('metatool.query', {'source': 'PubMed', 'q': query}) or []
         except Exception:
             pass
-        # Deduplicate by id
-        unique = {paper["id"]: paper for paper in results}
+        # dedupe
+        unique = {p['id']: p for p in papers if 'id' in p}
         return list(unique.values())[:limit]
 
     def get_notebook_cells(self, paper_id: str):
         """
-        Retrieve code cells for reproducible example from the Python-run MCP server.
+        Retrieve reproducible code cells from Python-run MCP server.
         """
         try:
-            resp = self.runner.call("mcp-run-python.execute", {"paper_id": paper_id})
-            return resp.get("cells", [])
+            result = self.runner.call('mcp-run-python.execute', {'paper_id': paper_id})
+            return result.get('cells', [])
         except Exception:
             return []
 
     def get_graph(self, paper_id: str):
         """
-        Retrieve a knowledge graph representation for a paper from the Chroma MCP server.
+        Retrieve a knowledge graph for a paper from the Chroma MCP server.
         """
         try:
-            return self.chroma.call("chroma.query_graph", {"id": paper_id})
+            return self.chroma.call('chroma.query_graph', {'id': paper_id}) or {'nodes': [], 'edges': []}
         except Exception:
-            return {"nodes": [], "edges": []}
+            return {'nodes': [], 'edges': []}