update

scripts/start_server.sh

@@ -0,0 +1 @@
+python -m uvicorn server:app --reload

server.py

@@ -23,38 +23,44 @@
     "storage": None,
     "memory": None,
     "tool": None,
-    "scheduler": None
+    "scheduler": None,
 }
 
 send_request, SysCallWrapper = useSysCall()
 
+
 class LLMConfig(BaseModel):
     llm_name: str
-    max_gpu_memory: dict | None = None 
+    max_gpu_memory: dict | None = None
     eval_device: str = "cuda:0"
     max_new_tokens: int = 2048
     log_mode: str = "INFO"
     use_backend: str = "default"
 
+
 class StorageConfig(BaseModel):
     root_dir: str = "root"
     use_vector_db: bool = False
     vector_db_config: Optional[Dict[str, Any]] = None
 
+
 class MemoryConfig(BaseModel):
     memory_limit: int = 104857600  # 100MB in bytes
     eviction_k: int = 10
     custom_eviction_policy: Optional[str] = None
 
+
 class ToolManagerConfig(BaseModel):
     allowed_tools: Optional[list[str]] = None
     custom_tools: Optional[Dict[str, Any]] = None
 
+
 class SchedulerConfig(BaseModel):
     log_mode: str = "INFO"
     max_workers: int = 64
     custom_syscalls: Optional[Dict[str, Any]] = None
 
+
 @app.post("/core/llm/setup")
 async def setup_llm(config: LLMConfig):
     """Set up the LLM core component."""
@@ -74,100 +80,118 @@ async def setup_llm(config: LLMConfig):
 
     #     raise HTTPException(status_code=500, detail=f"Failed to initialize LLM core: {str(e)}")
     llm = useCore(
-            llm_name=config.llm_name,
-            max_gpu_memory=config.max_gpu_memory,
-            eval_device=config.eval_device,
-            max_new_tokens=config.max_new_tokens,
-            log_mode=config.log_mode,
-            use_backend=config.use_backend
-        )
+        llm_name=config.llm_name,
+        max_gpu_memory=config.max_gpu_memory,
+        eval_device=config.eval_device,
+        max_new_tokens=config.max_new_tokens,
+        log_mode=config.log_mode,
+        use_backend=config.use_backend,
+    )
     active_components["llm"] = llm
     return {"status": "success", "message": "LLM core initialized"}
 
+
 @app.post("/core/storage/setup")
 async def setup_storage(config: StorageConfig):
     """Set up the storage manager component."""
     try:
         storage_manager = useStorageManager(
             root_dir=config.root_dir,
             use_vector_db=config.use_vector_db,
-            **(config.vector_db_config or {})
+            **(config.vector_db_config or {}),
         )
         active_components["storage"] = storage_manager
         return {"status": "success", "message": "Storage manager initialized"}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to initialize storage manager: {str(e)}")
+        raise HTTPException(
+            status_code=500, detail=f"Failed to initialize storage manager: {str(e)}"
+        )
+
 
 @app.post("/core/memory/setup")
 async def setup_memory(config: MemoryConfig):
     """Set up the memory manager component."""
     if not active_components["storage"]:
-        raise HTTPException(status_code=400, detail="Storage manager must be initialized first")
-
+        raise HTTPException(
+            status_code=400, detail="Storage manager must be initialized first"
+        )
+
     try:
         memory_manager = useMemoryManager(
             memory_limit=config.memory_limit,
             eviction_k=config.eviction_k,
-            storage_manager=active_components["storage"]
+            storage_manager=active_components["storage"],
         )
         active_components["memory"] = memory_manager
         return {"status": "success", "message": "Memory manager initialized"}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to initialize memory manager: {str(e)}")
+        raise HTTPException(
+            status_code=500, detail=f"Failed to initialize memory manager: {str(e)}"
+        )
+
 
 @app.post("/core/tool/setup")
 async def setup_tool_manager(config: ToolManagerConfig):
     """Set up the tool manager component."""
     try:
         tool_manager = useToolManager()
-        
+
         active_components["tool"] = tool_manager
         return {"status": "success", "message": "Tool manager initialized"}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to initialize tool manager: {str(e)}")
+        raise HTTPException(
+            status_code=500, detail=f"Failed to initialize tool manager: {str(e)}"
+        )
+
 
 @app.post("/core/factory/setup")
 async def setup_agent_factory(config: SchedulerConfig):
     """Set up the agent factory for managing agent execution."""
     required_components = ["llm", "memory", "storage", "tool"]
-    missing_components = [comp for comp in required_components if not active_components[comp]]
-
+    missing_components = [
+        comp for comp in required_components if not active_components[comp]
+    ]
+
     if missing_components:
         raise HTTPException(
             status_code=400,
-            detail=f"Missing required components: {', '.join(missing_components)}"
+            detail=f"Missing required components: {', '.join(missing_components)}",
         )
-    
+
     try:
         submit_agent, await_agent_execution = useFactory(
-            log_mode=config.log_mode,
-            max_workers=config.max_workers
+            log_mode=config.log_mode, max_workers=config.max_workers
         )
-        
+
         active_components["factory"] = {
             "submit": submit_agent,
-            "await": await_agent_execution
+            "await": await_agent_execution,
         }
 
-        print(active_components['llm'].model)
+        print(active_components["llm"].model)
 
         return {"status": "success", "message": "Agent factory initialized"}
     except Exception as e:
         print(e)
-        raise HTTPException(status_code=500, detail=f"Failed to initialize agent factory: {str(e)}")
+        raise HTTPException(
+            status_code=500, detail=f"Failed to initialize agent factory: {str(e)}"
+        )
+
 
 @app.post("/core/scheduler/setup")
 async def setup_scheduler(config: SchedulerConfig):
     """Set up the FIFO scheduler with all components."""
     required_components = ["llm", "memory", "storage", "tool"]
-    missing_components = [comp for comp in required_components if not active_components[comp]]
-
+    missing_components = [
+        comp for comp in required_components if not active_components[comp]
+    ]
+
     if missing_components:
         raise HTTPException(
             status_code=400,
-            detail=f"Missing required components: {', '.join(missing_components)}"
+            detail=f"Missing required components: {', '.join(missing_components)}",
         )
-    
+
     try:
         # Set up the scheduler with all components
         scheduler = fifo_scheduler(
@@ -179,17 +203,20 @@ async def setup_scheduler(config: SchedulerConfig):
             get_llm_syscall=None,
             get_memory_syscall=None,
             get_storage_syscall=None,
-            get_tool_syscall=None
+            get_tool_syscall=None,
             # **(config.custom_syscalls or {})
         )
-        
+
         active_components["scheduler"] = scheduler
 
         scheduler.start()
 
         return {"status": "success", "message": "Scheduler initialized"}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to initialize scheduler: {str(e)}")
+        raise HTTPException(
+            status_code=500, detail=f"Failed to initialize scheduler: {str(e)}"
+        )
+
 
 @app.get("/core/status")
 async def get_status():
@@ -199,59 +226,50 @@ async def get_status():
         for component, instance in active_components.items()
     }
 
+
 class AgentSubmit(BaseModel):
     agent_id: str
     agent_config: Dict[str, Any]
 
+
 @app.post("/agents/submit")
 async def submit_agent(config: AgentSubmit):
     """Submit an agent for execution using the agent factory."""
     if "factory" not in active_components or not active_components["factory"]:
-        raise HTTPException(
-            status_code=400,
-            detail="Agent factory not initialized"
-        )
-
+        raise HTTPException(status_code=400, detail="Agent factory not initialized")
+
     try:
         _submit_agent = active_components["factory"]["submit"]
-        execution_id = _submit_agent(agent_name=config.agent_id, task_input=config.agent_config['task'])
+        execution_id = _submit_agent(
+            agent_name=config.agent_id, task_input=config.agent_config["task"]
+        )
         print(execution_id)
-        
+
         return {
             "status": "success",
             "execution_id": execution_id,
-            "message": f"Agent {config.agent_id} submitted for execution"
+            "message": f"Agent {config.agent_id} submitted for execution",
         }
     except Exception as e:
         print(e)
-        raise HTTPException(
-            status_code=500,
-            detail=f"Failed to submit agent: {str(e)}"
-        )
+        raise HTTPException(status_code=500, detail=f"Failed to submit agent: {str(e)}")
+
 
 @app.get("/agents/{execution_id}/status")
 async def get_agent_status(execution_id: int):
     """Get the status of a submitted agent."""
     if "factory" not in active_components or not active_components["factory"]:
-        raise HTTPException(
-            status_code=400,
-            detail="Agent factory not initialized"
-        )
-
+        raise HTTPException(status_code=400, detail="Agent factory not initialized")
+
     try:
         await_execution = active_components["factory"]["await"]
         result = await_execution(int(execution_id))
 
-        return {
-            "status": "completed",
-            "result": result
-        }
+        return {"status": "completed", "result": result}
     except Exception as e:
         print(e)
-        return {
-            "status": "running",
-            "message": str(e)
-        }
+        return {"status": "running", "message": str(e)}
+
 
 @app.post("/core/cleanup")
 async def cleanup_components():
@@ -267,20 +285,20 @@ async def cleanup_components():
                     active_components[component].cleanup()
                 active_components[component] = None
 
-
-
-
-
         return {"status": "success", "message": "All components cleaned up"}
     except Exception as e:
         print(e)
-        raise HTTPException(status_code=500, detail=f"Failed to cleanup components: {str(e)}")
+        raise HTTPException(
+            status_code=500, detail=f"Failed to cleanup components: {str(e)}"
+        )
+
 
 class QueryRequest(BaseModel):
     agent_name: str
     query_type: Literal["llm", "tool", "storage", "memory"]
     query_data: LLMQuery
 
+
 @app.post("/query")
 async def handle_query(request: QueryRequest):
     try:
@@ -289,11 +307,12 @@ async def handle_query(request: QueryRequest):
                 messages=request.query_data.messages,
                 tools=request.query_data.tools,
                 action_type=request.query_data.action_type,
-                message_return_type=request.query_data.message_return_type
+                message_return_type=request.query_data.message_return_type,
             )
             return send_request(request.agent_name, query)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)