fixing rag

backend/tests/clients/chatbot_api_client.py

@@ -0,0 +1,355 @@
+"""
+Chatbot API test client for comprehensive workflow testing.
+"""
+
+import aiohttp
+import asyncio
+from typing import Dict, List, Optional, Any, AsyncGenerator
+import json
+import time
+from pathlib import Path
+
+
+class ChatbotAPITestClient:
+    """Test client for chatbot API workflows"""
+    
+    def __init__(self, base_url: str = "http://localhost:3001"):
+        self.base_url = base_url.rstrip('/')
+        self.session_timeout = aiohttp.ClientTimeout(total=60)
+        self.auth_token = None
+        self.api_key = None
+    
+    async def authenticate(self, email: str = "test@example.com", password: str = "testpass123") -> Dict[str, Any]:
+        """Authenticate user and get JWT token"""
+        login_data = {"email": email, "password": password}
+        
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.base_url}/api-internal/v1/auth/login",
+                json=login_data
+            ) as response:
+                if response.status == 200:
+                    result = await response.json()
+                    self.auth_token = result.get("access_token")
+                    return {"success": True, "token": self.auth_token}
+                else:
+                    error = await response.text()
+                    return {"success": False, "error": error, "status": response.status}
+    
+    async def register_user(self, email: str, password: str, username: str) -> Dict[str, Any]:
+        """Register a new user"""
+        user_data = {
+            "email": email,
+            "password": password,
+            "username": username
+        }
+        
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.base_url}/api-internal/v1/auth/register",
+                json=user_data
+            ) as response:
+                result = await response.json() if response.content_type == 'application/json' else await response.text()
+                return {
+                    "status_code": response.status,
+                    "success": response.status == 201,
+                    "data": result
+                }
+    
+    async def create_rag_collection(self, name: str, description: str = "") -> Dict[str, Any]:
+        """Create a RAG collection"""
+        if not self.auth_token:
+            return {"error": "Not authenticated"}
+        
+        collection_data = {
+            "name": name,
+            "description": description,
+            "processing_config": {
+                "chunk_size": 1000,
+                "chunk_overlap": 200
+            }
+        }
+        
+        headers = {"Authorization": f"Bearer {self.auth_token}"}
+        
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.base_url}/api-internal/v1/rag/collections",
+                json=collection_data,
+                headers=headers
+            ) as response:
+                result = await response.json() if response.content_type == 'application/json' else await response.text()
+                return {
+                    "status_code": response.status,
+                    "success": response.status == 201,
+                    "data": result
+                }
+    
+    async def upload_document(self, collection_id: str, file_content: str, filename: str) -> Dict[str, Any]:
+        """Upload document to RAG collection"""
+        if not self.auth_token:
+            return {"error": "Not authenticated"}
+        
+        headers = {"Authorization": f"Bearer {self.auth_token}"}
+        
+        # Create form data
+        data = aiohttp.FormData()
+        data.add_field('file', file_content, filename=filename, content_type='text/plain')
+        
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.base_url}/api-internal/v1/rag/collections/{collection_id}/upload",
+                data=data,
+                headers=headers
+            ) as response:
+                result = await response.json() if response.content_type == 'application/json' else await response.text()
+                return {
+                    "status_code": response.status,
+                    "success": response.status == 201,
+                    "data": result
+                }
+    
+    async def wait_for_document_processing(self, document_id: str, timeout: int = 60) -> Dict[str, Any]:
+        """Wait for document processing to complete"""
+        if not self.auth_token:
+            return {"error": "Not authenticated"}
+        
+        headers = {"Authorization": f"Bearer {self.auth_token}"}
+        start_time = time.time()
+        
+        async with aiohttp.ClientSession() as session:
+            while time.time() - start_time < timeout:
+                async with session.get(
+                    f"{self.base_url}/api-internal/v1/rag/documents/{document_id}",
+                    headers=headers
+                ) as response:
+                    if response.status == 200:
+                        doc = await response.json()
+                        status = doc.get("processing_status")
+                        
+                        if status == "completed":
+                            return {"success": True, "status": "completed", "document": doc}
+                        elif status == "failed":
+                            return {"success": False, "status": "failed", "document": doc}
+                        
+                        await asyncio.sleep(1)
+                    else:
+                        return {"success": False, "error": f"Failed to check status: {response.status}"}
+        
+        return {"success": False, "error": "Timeout waiting for processing"}
+    
+    async def create_chatbot(self, 
+                           name: str,
+                           chatbot_type: str = "assistant",
+                           use_rag: bool = False,
+                           rag_collection: str = None,
+                           **config) -> Dict[str, Any]:
+        """Create a chatbot"""
+        if not self.auth_token:
+            return {"error": "Not authenticated"}
+        
+        chatbot_data = {
+            "name": name,
+            "chatbot_type": chatbot_type,
+            "model": "test-model",
+            "system_prompt": "You are a helpful assistant.",
+            "use_rag": use_rag,
+            "rag_collection": rag_collection,
+            "rag_top_k": 3,
+            "temperature": 0.7,
+            "max_tokens": 1000,
+            **config
+        }
+        
+        headers = {"Authorization": f"Bearer {self.auth_token}"}
+        
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.base_url}/api-internal/v1/chatbot/create",
+                json=chatbot_data,
+                headers=headers
+            ) as response:
+                result = await response.json() if response.content_type == 'application/json' else await response.text()
+                return {
+                    "status_code": response.status,
+                    "success": response.status == 201,
+                    "data": result
+                }
+    
+    async def create_api_key_for_chatbot(self, chatbot_id: str, name: str = "Test API Key") -> Dict[str, Any]:
+        """Create API key for chatbot"""
+        if not self.auth_token:
+            return {"error": "Not authenticated"}
+        
+        api_key_data = {
+            "name": name,
+            "scopes": ["chatbot.chat"],
+            "budget_limit": 100.0,
+            "chatbot_id": chatbot_id
+        }
+        
+        headers = {"Authorization": f"Bearer {self.auth_token}"}
+        
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.base_url}/api-internal/v1/chatbot/{chatbot_id}/api-key",
+                json=api_key_data,
+                headers=headers
+            ) as response:
+                result = await response.json() if response.content_type == 'application/json' else await response.text()
+                if response.status == 201 and isinstance(result, dict):
+                    self.api_key = result.get("key")
+                return {
+                    "status_code": response.status,
+                    "success": response.status == 201,
+                    "data": result
+                }
+    
+    async def chat_with_bot(self, 
+                          chatbot_id: str, 
+                          message: str, 
+                          conversation_id: Optional[str] = None,
+                          api_key: Optional[str] = None) -> Dict[str, Any]:
+        """Send message to chatbot"""
+        if not api_key and not self.api_key:
+            return {"error": "No API key available"}
+        
+        chat_data = {
+            "message": message,
+            "conversation_id": conversation_id
+        }
+        
+        headers = {"Authorization": f"Bearer {api_key or self.api_key}"}
+        
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.base_url}/api/v1/chatbot/{chatbot_id}/chat",
+                json=chat_data,
+                headers=headers
+            ) as response:
+                result = await response.json() if response.content_type == 'application/json' else await response.text()
+                return {
+                    "status_code": response.status,
+                    "success": response.status == 200,
+                    "data": result
+                }
+    
+    async def test_rag_workflow(self,
+                              collection_name: str,
+                              document_content: str,
+                              chatbot_name: str,
+                              test_question: str) -> Dict[str, Any]:
+        """Test complete RAG workflow from document upload to chat"""
+        workflow_results = {}
+        
+        # Step 1: Create RAG collection
+        collection_result = await self.create_rag_collection(collection_name, "Test collection for workflow")
+        workflow_results["collection_creation"] = collection_result
+        
+        if not collection_result["success"]:
+            return {"success": False, "error": "Failed to create collection", "results": workflow_results}
+        
+        collection_id = collection_result["data"]["id"]
+        
+        # Step 2: Upload document
+        document_result = await self.upload_document(collection_id, document_content, "test_doc.txt")
+        workflow_results["document_upload"] = document_result
+        
+        if not document_result["success"]:
+            return {"success": False, "error": "Failed to upload document", "results": workflow_results}
+        
+        document_id = document_result["data"]["id"]
+        
+        # Step 3: Wait for processing
+        processing_result = await self.wait_for_document_processing(document_id)
+        workflow_results["document_processing"] = processing_result
+        
+        if not processing_result["success"]:
+            return {"success": False, "error": "Document processing failed", "results": workflow_results}
+        
+        # Step 4: Create chatbot with RAG
+        chatbot_result = await self.create_chatbot(
+            name=chatbot_name,
+            use_rag=True,
+            rag_collection=collection_name
+        )
+        workflow_results["chatbot_creation"] = chatbot_result
+        
+        if not chatbot_result["success"]:
+            return {"success": False, "error": "Failed to create chatbot", "results": workflow_results}
+        
+        chatbot_id = chatbot_result["data"]["id"]
+        
+        # Step 5: Create API key
+        api_key_result = await self.create_api_key_for_chatbot(chatbot_id)
+        workflow_results["api_key_creation"] = api_key_result
+        
+        if not api_key_result["success"]:
+            return {"success": False, "error": "Failed to create API key", "results": workflow_results}
+        
+        # Step 6: Test chat with RAG
+        chat_result = await self.chat_with_bot(chatbot_id, test_question)
+        workflow_results["chat_test"] = chat_result
+        
+        if not chat_result["success"]:
+            return {"success": False, "error": "Chat test failed", "results": workflow_results}
+        
+        # Step 7: Verify RAG sources in response
+        chat_response = chat_result["data"]
+        has_sources = "sources" in chat_response and len(chat_response["sources"]) > 0
+        workflow_results["rag_verification"] = {
+            "has_sources": has_sources,
+            "source_count": len(chat_response.get("sources", [])),
+            "sources": chat_response.get("sources", [])
+        }
+        
+        return {
+            "success": True,
+            "workflow_complete": True,
+            "rag_working": has_sources,
+            "results": workflow_results
+        }
+    
+    async def test_conversation_memory(self, chatbot_id: str, api_key: str = None) -> Dict[str, Any]:
+        """Test conversation memory functionality"""
+        messages = [
+            "My name is Alice and I like cats.",
+            "What's my name?",
+            "What do I like?"
+        ]
+        
+        conversation_id = None
+        results = []
+        
+        for i, message in enumerate(messages):
+            result = await self.chat_with_bot(chatbot_id, message, conversation_id, api_key)
+            
+            if result["success"]:
+                conversation_id = result["data"].get("conversation_id")
+                results.append({
+                    "message_index": i,
+                    "message": message,
+                    "response": result["data"].get("response"),
+                    "conversation_id": conversation_id
+                })
+            else:
+                results.append({
+                    "message_index": i,
+                    "message": message,
+                    "error": result.get("error"),
+                    "status_code": result.get("status_code")
+                })
+        
+        # Analyze memory performance
+        memory_working = False
+        if len(results) >= 3:
+            # Check if the bot remembers the name in the second response
+            response2 = results[1].get("response", "").lower()
+            response3 = results[2].get("response", "").lower()
+            memory_working = "alice" in response2 and ("cat" in response3 or "like" in response3)
+        
+        return {
+            "conversation_results": results,
+            "memory_working": memory_working,
+            "conversation_maintained": all(r.get("conversation_id") == results[0].get("conversation_id") for r in results if r.get("conversation_id"))
+        }