v0 of user rolled AI

vocode/user_implemented_agent/base_agent.py

@@ -9,5 +9,5 @@ class BaseAgent():
     async def respond(self, human_input) -> str:
         raise NotImplementedError
     
-    def run(self, host="localhost", port=3001):
+    def run(self, host="localhost", port=3000):
         uvicorn.run(self.app, host=host, port=port)

vocode/user_implemented_agent/restful_agent.py

@@ -1,17 +1,15 @@
 from .base_agent import BaseAgent
+from ..models.agent import RESTfulAgentInput, RESTfulAgentOutput
 from pydantic import BaseModel
 from fastapi import APIRouter
 
 class RESTfulAgent(BaseAgent):
-
-    class HumanInput(BaseModel):
-        human_input: str
         
     def __init__(self):
         super().__init__()
         self.app.post("/respond")(self.respond_rest)
 
-    async def respond_rest(self, request: HumanInput):
+    async def respond_rest(self, request: RESTfulAgentInput) -> RESTfulAgentOutput:
         response = await self.respond(request.human_input)
-        return {"response": response}
+        return RESTfulAgentOutput(response=response)
 

vocode/user_implemented_agent/websocket_agent.py

@@ -0,0 +1,26 @@
+from .base_agent import BaseAgent
+from pydantic import BaseModel
+import typing
+from fastapi import APIRouter, WebSocket
+from ..models.agent import AgentStartMessage, AgentReadyMessage, AgentTextMessage, WebSocketAgentMessage, WebSocketAgentMessageType
+from jsonpath_ng import parse
+
+class WebSocketAgent(BaseAgent):
+        
+    def __init__(self):
+        super().__init__()
+        self.app.websocket("/respond")(self.respond_websocket)
+
+    async def respond_websocket(self, websocket: WebSocket):
+        await websocket.accept()
+        AgentStartMessage.parse_obj(await websocket.receive_json())
+        await websocket.send_text(AgentReadyMessage().json())
+        while True:
+            message = WebSocketAgentMessage.parse_obj(await websocket.receive_json())
+            if message.type == WebSocketAgentMessageType.AGENT_STOP:
+                break
+            text_message = typing.cast(AgentTextMessage, message)
+            response = await self.respond(text_message.data.text)
+            await websocket.send_text(AgentTextMessage.from_text(response).json())
+        await websocket.close()
+