Merge pull request #49 from EnvisionMindCa/codex/create-api-for-chat-functionality

README.md

@@ -121,3 +121,24 @@ python run.py
 
 The custom VM includes typical utilities like ``sudo`` and ``curl`` so it behaves
 more like a standard Ubuntu installation.
+
+## REST API
+
+Start the API server using ``uvicorn``:
+
+```bash
+uvicorn src.api:app --host 0.0.0.0 --port 8000
+```
+
+### Endpoints
+
+- ``POST /chat/stream`` – Stream the assistant's response as plain text.
+- ``POST /upload`` – Upload a document so it can be referenced in chats.
+
+Example request:
+
+```bash
+curl -N -X POST http://localhost:8000/chat/stream \
+     -H 'Content-Type: application/json' \
+     -d '{"user":"demo","session":"default","prompt":"Hello"}'
+```

requirements.txt

@@ -6,3 +6,4 @@ colorama
 python-dotenv
 fastapi
 uvicorn
+python-multipart

run.py

@@ -10,7 +10,6 @@ async def _main() -> None:
     async with ChatSession(user="demo_user", session="demo_session") as chat:
         doc_path = chat.upload_document("note.pdf")
         # print(f"Document uploaded to VM at: {doc_path}")
-        # answer = await chat.chat(f"Remove all contents of test.txt and add the text 'Hello, World!' to it.")
         # async for resp in chat.chat_stream("Erase the contents of test.txt and write 'Hello, World!' to it."):
         # async for resp in chat.chat_stream("Verify that the file test.txt exists and contains the text 'Hello, World!'."):
         # async for resp in chat.chat_stream("Inspect the contents of note.pdf and summarize it."):

src/api.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from fastapi import FastAPI, UploadFile, File, Form
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+import asyncio
+import os
+import tempfile
+from pathlib import Path
+
+from .chat import ChatSession
+from .log import get_logger
+
+
+_LOG = get_logger(__name__)
+
+
+class ChatRequest(BaseModel):
+    user: str = "default"
+    session: str = "default"
+    prompt: str
+
+
+def create_app() -> FastAPI:
+    app = FastAPI(title="LLM Backend API")
+
+    @app.post("/chat/stream")
+    async def chat_stream(req: ChatRequest):
+        async def stream() -> asyncio.AsyncIterator[str]:
+            async with ChatSession(user=req.user, session=req.session) as chat:
+                try:
+                    async for part in chat.chat_stream(req.prompt):
+                        yield part
+                except Exception as exc:  # pragma: no cover - runtime failures
+                    _LOG.error("Streaming chat failed: %s", exc)
+                    yield f"Error: {exc}"
+
+        return StreamingResponse(stream(), media_type="text/plain")
+
+    @app.post("/upload")
+    async def upload_document(
+        user: str = Form(...),
+        session: str = Form("default"),
+        file: UploadFile = File(...),
+    ):
+        async with ChatSession(user=user, session=session) as chat:
+            tmpdir = tempfile.mkdtemp(prefix="upload_")
+            tmp_path = Path(tmpdir) / file.filename
+            try:
+                contents = await file.read()
+                tmp_path.write_bytes(contents)
+                vm_path = chat.upload_document(str(tmp_path))
+            finally:
+                try:
+                    os.remove(tmp_path)
+                    os.rmdir(tmpdir)
+                except OSError:
+                    pass
+        return {"path": vm_path}
+
+    @app.get("/health")
+    async def health():
+        return {"status": "ok"}
+
+    return app
+
+
+app = create_app()
+
+if __name__ == "__main__":  # pragma: no cover - manual start
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "8000")))

src/chat.py

@@ -294,33 +294,6 @@ class ChatSession:
         async with self._lock:
             self._state = "idle"
 
-    async def _handle_tool_calls(
-        self,
-        messages: List[Msg],
-        response: ChatResponse,
-        conversation: Conversation,
-        depth: int = 0,
-    ) -> ChatResponse:
-        final = response
-        gen = self._handle_tool_calls_stream(messages, response, conversation, depth)
-        async for final in gen:
-            pass
-        return final
-
-    async def chat(self, prompt: str) -> str:
-        DBMessage.create(conversation=self._conversation, role="user", content=prompt)
-        self._messages.append({"role": "user", "content": prompt})
-
-        response = await self.ask(self._messages)
-        self._messages.append(response.message.model_dump())
-        self._store_assistant_message(self._conversation, response.message)
-
-        _LOG.info("Thinking:\n%s", response.message.thinking or "<no thinking trace>")
-
-        final_resp = await self._handle_tool_calls(
-            self._messages, response, self._conversation
-        )
-        return self._format_output(final_resp.message)
 
     async def chat_stream(self, prompt: str) -> AsyncIterator[str]:
         async with self._lock: