Integrate TxAgent into CPS API: add TxAgent files, update requirements, fix imports and database collections

api/__init__.py

@@ -17,7 +17,11 @@ def get_router():
     from .routes.pdf import pdf
     from .routes.appointments import router as appointments
     from .routes.messaging import router as messaging
-    from .routes.txagent import router as txagent
+    from .routes.txagent import router as txagent, get_txagent, ChatRequest
+    from fastapi import Depends, HTTPException
+    from core.security import get_current_user
+    import asyncio
+    from fastapi.responses import StreamingResponse
     
     # Include sub-routers with correct prefixes
     router.include_router(patients, tags=["patients"])  # Remove prefix since routes already have /patients
@@ -27,6 +31,65 @@ def get_router():
     router.include_router(messaging, tags=["messaging"])
     router.include_router(txagent, tags=["txagent"])
     
+    # Add the chat-stream endpoint at the root level
+    @router.post("/chat-stream")
+    async def chat_stream(
+        request: ChatRequest,
+        current_user: dict = Depends(get_current_user)
+    ):
+        """Streaming chat endpoint for TxAgent"""
+        import logging
+        logger = logging.getLogger(__name__)
+        
+        try:
+            if not any(role in current_user.get('roles', []) for role in ['doctor', 'admin']):
+                raise HTTPException(status_code=403, detail="Only doctors and admins can use TxAgent")
+
+            # Get TxAgent instance
+            txagent = get_txagent()
+            
+            if txagent is None:
+                # Fallback response if TxAgent is not available
+                response_text = f"TxAgent is currently unavailable. Please try again later. Your message was: {request.message}"
+            else:
+                # Get real response from TxAgent
+                try:
+                    response_text = txagent.chat(
+                        message=request.message,
+                        temperature=request.temperature,
+                        max_new_tokens=request.max_new_tokens
+                    )
+                except Exception as e:
+                    logger.error(f"TxAgent chat failed: {e}")
+                    response_text = f"Sorry, I encountered an error processing your request: {request.message}. Please try again."
+
+            # Create streaming response
+            async def generate():
+                # Split response into sentences for more natural streaming
+                sentences = response_text.split('. ')
+                
+                for i, sentence in enumerate(sentences):
+                    if sentence.strip():
+                        # Add period back except for the last sentence
+                        chunk = sentence + ('. ' if i < len(sentences) - 1 else '')
+                        yield f"data: {chunk}\n\n"
+                        await asyncio.sleep(0.2)  # Natural delay between sentences
+                
+                yield "data: [DONE]\n\n"
+
+            return StreamingResponse(
+                generate(),
+                media_type="text/plain",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "Content-Type": "text/event-stream"
+                }
+            )
+        except Exception as e:
+            logger.error(f"Error in chat stream: {e}")
+            raise HTTPException(status_code=500, detail="Failed to process chat stream")
+    
     return router
 
 # Export the router

api/routes/txagent.py

@@ -265,57 +265,4 @@ async def get_chats(current_user: dict = Depends(get_current_user)):
         logger.error(f"Error getting chats: {e}")
         raise HTTPException(status_code=500, detail="Failed to get chats")
 
-@router.post("/chat-stream")
-async def chat_stream(
-    request: ChatRequest,
-    current_user: dict = Depends(get_current_user)
-):
-    """Streaming chat endpoint for TxAgent"""
-    try:
-        if not any(role in current_user.get('roles', []) for role in ['doctor', 'admin']):
-            raise HTTPException(status_code=403, detail="Only doctors and admins can use TxAgent")
 
-        # Get TxAgent instance
-        txagent = get_txagent()
-        
-        if txagent is None:
-            # Fallback response if TxAgent is not available
-            response_text = f"TxAgent is currently unavailable. Please try again later. Your message was: {request.message}"
-        else:
-            # Get real response from TxAgent
-            try:
-                response_text = txagent.chat(
-                    message=request.message,
-                    temperature=request.temperature,
-                    max_new_tokens=request.max_new_tokens
-                )
-            except Exception as e:
-                logger.error(f"TxAgent chat failed: {e}")
-                response_text = f"Sorry, I encountered an error processing your request: {request.message}. Please try again."
-
-        # Create streaming response
-        async def generate():
-            # Split response into sentences for more natural streaming
-            sentences = response_text.split('. ')
-            
-            for i, sentence in enumerate(sentences):
-                if sentence.strip():
-                    # Add period back except for the last sentence
-                    chunk = sentence + ('. ' if i < len(sentences) - 1 else '')
-                    yield f"data: {chunk}\n\n"
-                    await asyncio.sleep(0.2)  # Natural delay between sentences
-            
-            yield "data: [DONE]\n\n"
-
-        return StreamingResponse(
-            generate(),
-            media_type="text/plain",
-            headers={
-                "Cache-Control": "no-cache",
-                "Connection": "keep-alive",
-                "Content-Type": "text/event-stream"
-            }
-        )
-    except Exception as e:
-        logger.error(f"Error in chat stream: {e}")
-        raise HTTPException(status_code=500, detail="Failed to process chat stream")