Spaces:
Running
Running
| # from fastapi import FastAPI, HTTPException, Query as QueryParam | |
| # from pydantic import BaseModel, Field | |
| # from langchain_openai import ChatOpenAI, OpenAIEmbeddings | |
| # from qdrant_client import QdrantClient | |
| # from langchain.agents import Tool, AgentExecutor, create_openai_tools_agent | |
| # from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder | |
| # from langchain.memory import ConversationBufferMemory | |
| # from typing import Optional, List, Dict, Any | |
| # import os | |
| # import warnings | |
| # import base64 | |
| # import requests | |
| # from dotenv import load_dotenv | |
| # from datetime import datetime | |
| # import json | |
| # import uuid | |
| # import redis | |
| # # Suppress warnings | |
| # warnings.filterwarnings("ignore", message="Qdrant client version.*is incompatible.*") | |
| # load_dotenv() | |
| # app = FastAPI(title="AI Agent with Redis Session Management") | |
| # # Environment variables | |
| # OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") | |
| # QDRANT_COLLECTION_NAME = os.getenv("QDRANT_COLLECTION_NAME", "vatsav_test_1") | |
| # QDRANT_HOST = os.getenv("QDRANT_HOST", "127.0.0.1") | |
| # QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333)) | |
| # # Redis Configuration | |
| # REDIS_URL = os.getenv("REDIS_URL") | |
| # REDIS_HOST = os.getenv("REDIS_HOST", "127.0.0.1") | |
| # REDIS_PORT = int(os.getenv("REDIS_PORT", 6379)) | |
| # REDIS_PASSWORD = os.getenv("REDIS_PASSWORD") | |
| # # Initialize Redis client | |
| # def get_redis_client(): | |
| # """Initialize Redis client with fallback to local Redis""" | |
| # try: | |
| # if REDIS_URL: | |
| # # Use deployed Redis URL | |
| # redis_client = redis.from_url( | |
| # REDIS_URL, | |
| # decode_responses=True, | |
| # socket_connect_timeout=5, | |
| # socket_timeout=5 | |
| # ) | |
| # # Test connection | |
| # redis_client.ping() | |
| # print(f"✅ Connected to deployed Redis: {REDIS_URL}") | |
| # return redis_client | |
| # else: | |
| # # Use local Redis | |
| # redis_client = redis.StrictRedis( | |
| # host=REDIS_HOST, | |
| # port=REDIS_PORT, | |
| # password=REDIS_PASSWORD, | |
| # decode_responses=True, | |
| # socket_connect_timeout=5, | |
| # socket_timeout=5 | |
| # ) | |
| # # Test connection | |
| # redis_client.ping() | |
| # print(f"✅ Connected to local Redis: {REDIS_HOST}:{REDIS_PORT}") | |
| # return redis_client | |
| # except Exception as e: | |
| # print(f"❌ Redis connection failed: {e}") | |
| # raise HTTPException(status_code=500, detail=f"Redis connection failed: {str(e)}") | |
| # # Initialize Redis client | |
| # redis_client = get_redis_client() | |
| # # Initialize models | |
| # embedding_model = OpenAIEmbeddings( | |
| # model="text-embedding-3-large", | |
| # openai_api_key=OPENAI_API_KEY, | |
| # ) | |
| # qdrant_client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT) | |
| # llm = ChatOpenAI(model="gpt-4o", temperature=0, openai_api_key=OPENAI_API_KEY) | |
| # # === INPUT SCHEMAS === | |
| # class Query(BaseModel): | |
| # message: str | |
| # class ProjectRequest(BaseModel): | |
| # userLoginId: int | |
| # orgId: int | |
| # auth_token: str | |
| # class BotQuery(BaseModel): | |
| # userLoginId: int | |
| # orgId: int | |
| # auth_token: str | |
| # session_id: Optional[str] = None | |
| # message: str | |
| # class SessionResponse(BaseModel): | |
| # session_id: str | |
| # userLoginId: int | |
| # orgId: int | |
| # created_at: str | |
| # status: str | |
| # class MessageResponse(BaseModel): | |
| # message_id: str | |
| # session_id: str | |
| # role: str # "user" or "assistant" | |
| # message: str | |
| # timestamp: str | |
| # class ChatHistoryResponse(BaseModel): | |
| # session_id: str | |
| # messages: List[MessageResponse] | |
| # total_messages: int | |
| # # === SESSION MANAGEMENT FUNCTIONS === | |
| # def create_session(userLoginId: int, orgId: int, auth_token: str) -> dict: | |
| # """Create a new chat session""" | |
| # session_id = str(uuid.uuid4()) | |
| # session_data = { | |
| # "session_id": session_id, | |
| # "userLoginId": userLoginId, | |
| # "orgId": orgId, | |
| # "auth_token": auth_token, | |
| # "created_at": datetime.now().isoformat(), | |
| # "status": "active" | |
| # } | |
| # # Store session in Redis with 24 hour TTL | |
| # redis_client.setex( | |
| # f"session:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps(session_data) | |
| # ) | |
| # # Initialize empty chat history | |
| # redis_client.setex( | |
| # f"chat:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps([]) | |
| # ) | |
| # # Initialize conversation memory | |
| # redis_client.setex( | |
| # f"memory:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps([]) | |
| # ) | |
| # return session_data | |
| # def get_session(session_id: str) -> dict: | |
| # """Get session data from Redis""" | |
| # session_data = redis_client.get(f"session:{session_id}") | |
| # if not session_data: | |
| # raise HTTPException(status_code=404, detail="Session not found or expired") | |
| # return json.loads(session_data) | |
| # def add_message_to_session(session_id: str, role: str, message: str) -> str: | |
| # """Add message to session chat history""" | |
| # message_id = str(uuid.uuid4()) | |
| # message_data = { | |
| # "message_id": message_id, | |
| # "session_id": session_id, | |
| # "role": role, | |
| # "message": message, | |
| # "timestamp": datetime.now().isoformat() | |
| # } | |
| # # Get current chat history | |
| # chat_history = redis_client.get(f"chat:{session_id}") | |
| # if chat_history: | |
| # messages = json.loads(chat_history) | |
| # else: | |
| # messages = [] | |
| # # Add new message | |
| # messages.append(message_data) | |
| # # Update chat history in Redis with extended TTL | |
| # redis_client.setex( | |
| # f"chat:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps(messages) | |
| # ) | |
| # return message_id | |
| # def get_session_memory(session_id: str) -> List[Dict]: | |
| # """Get conversation memory for session""" | |
| # memory_data = redis_client.get(f"memory:{session_id}") | |
| # if memory_data: | |
| # return json.loads(memory_data) | |
| # return [] | |
| # def update_session_memory(session_id: str, messages: List[Dict]): | |
| # """Update conversation memory for session""" | |
| # redis_client.setex( | |
| # f"memory:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps(messages) | |
| # ) | |
| # def get_user_sessions(userLoginId: int) -> List[dict]: | |
| # """Get all sessions for a user""" | |
| # sessions = [] | |
| # # Scan for all session keys | |
| # for key in redis_client.scan_iter(match="session:*"): | |
| # session_data = redis_client.get(key) | |
| # if session_data: | |
| # session = json.loads(session_data) | |
| # if session["userLoginId"] == userLoginId: | |
| # sessions.append(session) | |
| # return sessions | |
| # def delete_session(session_id: str): | |
| # """Delete session and associated data""" | |
| # # Delete session data | |
| # redis_client.delete(f"session:{session_id}") | |
| # # Delete chat history | |
| # redis_client.delete(f"chat:{session_id}") | |
| # # Delete memory | |
| # redis_client.delete(f"memory:{session_id}") | |
| # # === UTILITY FUNCTIONS === | |
| # def get_encoded_auth_token(user: int, token: str) -> str: | |
| # auth_string = f"{user}:{token}" | |
| # return base64.b64encode(auth_string.encode("utf-8")).decode("utf-8") | |
| # def fetch_user_projects(userLoginId: int, orgId: int, auth_token: str): | |
| # url = "https://japidemo.dev.ingenspark.com/fetchUserProjects" | |
| # payload = { | |
| # "userLoginId": userLoginId, | |
| # "orgId": orgId | |
| # } | |
| # headers = { | |
| # 'accept': 'application/json, text/plain, */*', | |
| # 'authorization': f'Basic {auth_token}', | |
| # 'content-type': 'application/json; charset=UTF-8' | |
| # } | |
| # try: | |
| # response = requests.post(url, headers=headers, json=payload) | |
| # response.raise_for_status() | |
| # return response.json() | |
| # except requests.exceptions.RequestException as e: | |
| # raise HTTPException(status_code=response.status_code if 'response' in locals() else 500, | |
| # detail=str(e)) | |
| # def format_project_response(data: dict) -> str: | |
| # my_projects = data.get("data", {}).get("Myprojects", []) | |
| # other_projects = data.get("data", {}).get("Otherprojects", []) | |
| # all_projects = [] | |
| # for project in my_projects: | |
| # all_projects.append({ | |
| # "type": "Your Project", | |
| # "projectNm": project["projectNm"], | |
| # "projectId": project["projectId"], | |
| # "created_dttm": project["created_dttm"].split('.')[0], | |
| # "description": project["description"], | |
| # "categoryName": project["categoryName"] | |
| # }) | |
| # for project in other_projects: | |
| # all_projects.append({ | |
| # "type": "Other Project", | |
| # "projectNm": project["projectNm"], | |
| # "projectId": project["projectId"], | |
| # "created_dttm": project["created_dttm"].split('.')[0], | |
| # "description": project["description"], | |
| # "categoryName": project["categoryName"] | |
| # }) | |
| # if not all_projects: | |
| # return "❌ No projects found." | |
| # # Build the formatted string | |
| # result = [f"✅ You have access to {len(all_projects)} project(s):\n"] | |
| # for i, project in enumerate(all_projects, 1): | |
| # result.append(f"{i}. Project Name: {project['projectNm']} ({project['type']})") | |
| # result.append(f" Project ID: {project['projectId']}") | |
| # result.append(f" Created On: {project['created_dttm']}") | |
| # result.append(f" Description: {project['description']}") | |
| # result.append(f" Category: {project['categoryName']}\n") | |
| # return "\n".join(result) | |
| # # === TOOL FUNCTIONS === | |
| # def search_documents(query: str) -> str: | |
| # """Search through ingested documents and get relevant information.""" | |
| # try: | |
| # # Generate embedding for the query | |
| # query_vector = embedding_model.embed_query(query) | |
| # # Search in Qdrant | |
| # search_result = qdrant_client.search( | |
| # collection_name=QDRANT_COLLECTION_NAME, | |
| # query_vector=query_vector, | |
| # limit=5, | |
| # ) | |
| # if not search_result: | |
| # return "No relevant information found in the knowledge base." | |
| # # Convert results to text content | |
| # context_texts = [] | |
| # sources = [] | |
| # for hit in search_result: | |
| # context_texts.append(hit.payload["text"]) | |
| # sources.append(hit.payload.get("source", "unknown")) | |
| # # Create a simple prompt for answering based on context | |
| # context = "\n\n".join(context_texts) | |
| # unique_sources = list(set(sources)) | |
| # # Use the LLM directly to answer the message based on context | |
| # prompt = f"""Based on the following context, answer the message: {query} | |
| # Context: | |
| # {context} | |
| # Please provide a comprehensive answer based on the context above. If the context doesn't contain enough information to answer the message, say so clearly.""" | |
| # response = llm.invoke(prompt) | |
| # return f"{response.content}\n\nSources: {', '.join(unique_sources)}" | |
| # except Exception as e: | |
| # return f"Error searching documents: {str(e)}" | |
| # # Global variables to store auth context (for tool functions) | |
| # _current_user_id = None | |
| # _current_org_id = None | |
| # _current_auth_token = None | |
| # def get_user_projects(userLoginId: str) -> str: | |
| # """Get list of projects for a user.""" | |
| # try: | |
| # # Use global auth context if available | |
| # if _current_auth_token and _current_user_id: | |
| # user_id = _current_user_id | |
| # org_id = _current_org_id or 1 | |
| # auth_token = _current_auth_token | |
| # else: | |
| # return "❌ Authentication token required. Please provide auth_token in your request." | |
| # # Encode auth token using the actual user ID and provided token | |
| # encoded_token = get_encoded_auth_token(user_id, auth_token) | |
| # # Fetch projects | |
| # data = fetch_user_projects(user_id, org_id, encoded_token) | |
| # # Format and return the project list | |
| # formatted = format_project_response(data) | |
| # return formatted | |
| # except ValueError: | |
| # return "❌ Invalid userLoginId format. Please provide a valid number." | |
| # except Exception as e: | |
| # return f"❌ Error fetching projects: {str(e)}" | |
| # # === CREATE TOOLS === | |
| # document_search_tool = Tool( | |
| # name="document_search", | |
| # description="""Use this tool to search through ingested documents and get relevant information from the knowledge base. | |
| # Perfect for answering messages about uploaded documents, manuals, or any content that was previously stored. | |
| # Input should be a search query or message about the documents.""", | |
| # func=search_documents | |
| # ) | |
| # project_list_tool = Tool( | |
| # name="get_user_projects", | |
| # description="""Use this tool to get the list of projects for a user. | |
| # Perfect for when users ask about their projects, want to see available projects, or need project information. | |
| # Input should be the userLoginId (e.g., '25'). | |
| # Note: This tool requires authentication context to be set.""", | |
| # func=get_user_projects | |
| # ) | |
| # # === AGENT SETUP === | |
| # def create_agent_with_session_memory(session_id: str): | |
| # """Create agent with session memory from Redis""" | |
| # # Get memory from Redis | |
| # memory_messages = get_session_memory(session_id) | |
| # agent_prompt = ChatPromptTemplate.from_messages([ | |
| # ("system", """You are a helpful AI assistant with access to multiple tools and conversation memory: | |
| # 1. **Document Search**: Search through uploaded documents and knowledge base | |
| # 2. **Project Management**: Get list of user projects and project information | |
| # Your capabilities: | |
| # - Answer messages about documents using the document search tool | |
| # - Help users find their projects and project information | |
| # - Remember previous conversations in this session | |
| # - Provide general assistance and information | |
| # - Use appropriate tools based on user queries | |
| # Guidelines: | |
| # - Use the document search tool when users ask about specific content, documentation, or information that might be in uploaded files | |
| # - Use the project tool when users ask about projects, want to see their projects, or need project-related information | |
| # - Reference previous conversation context when relevant | |
| # - Be clear about which tool you're using and what information you're providing | |
| # - If you're unsure which tool to use, you can ask for clarification | |
| # - Provide helpful, accurate, and well-formatted responses | |
| # Remember: Always use the most appropriate tool based on the user's message and conversation context to provide the best possible answer."""), | |
| # MessagesPlaceholder(variable_name="chat_history"), | |
| # ("user", "{input}"), | |
| # MessagesPlaceholder(variable_name="agent_scratchpad"), | |
| # ]) | |
| # # Create memory object | |
| # memory = ConversationBufferMemory( | |
| # memory_key="chat_history", | |
| # return_messages=True | |
| # ) | |
| # # Load existing messages into memory | |
| # for msg in memory_messages: | |
| # if msg["role"] == "user": | |
| # memory.chat_memory.add_user_message(msg["message"]) | |
| # else: | |
| # memory.chat_memory.add_ai_message(msg["message"]) | |
| # # Create tools list | |
| # tools = [document_search_tool, project_list_tool] | |
| # # Create the agent | |
| # agent = create_openai_tools_agent(llm, tools, agent_prompt) | |
| # # Create the agent executor with memory | |
| # agent_executor = AgentExecutor( | |
| # agent=agent, | |
| # tools=tools, | |
| # verbose=True, | |
| # memory=memory | |
| # ) | |
| # return agent_executor, memory | |
| # # === API ENDPOINTS === | |
| # @app.post("/sessions", response_model=SessionResponse) | |
| # def create_new_session(userLoginId: int, orgId: int, auth_token: str): | |
| # """Create a new chat session""" | |
| # try: | |
| # session_data = create_session(userLoginId, orgId, auth_token) | |
| # return SessionResponse(**session_data) | |
| # except Exception as e: | |
| # raise HTTPException(status_code=500, detail=f"Error creating session: {str(e)}") | |
| # @app.get("/sessions") | |
| # def list_user_sessions(userLoginId: int): | |
| # """List all sessions for a user""" | |
| # try: | |
| # sessions = get_user_sessions(userLoginId) | |
| # return { | |
| # "userLoginId": userLoginId, | |
| # "total_sessions": len(sessions), | |
| # "sessions": sessions | |
| # } | |
| # except Exception as e: | |
| # raise HTTPException(status_code=500, detail=f"Error fetching sessions: {str(e)}") | |
| # @app.delete("/sessions/{session_id}") | |
| # def delete_user_session(session_id: str): | |
| # """Delete/close a session""" | |
| # try: | |
| # # Verify session exists | |
| # get_session(session_id) | |
| # # Delete session | |
| # delete_session(session_id) | |
| # return { | |
| # "message": f"Session {session_id} deleted successfully", | |
| # "session_id": session_id | |
| # } | |
| # except Exception as e: | |
| # raise HTTPException(status_code=500, detail=f"Error deleting session: {str(e)}") | |
| # @app.post("/bot") | |
| # def chat_with_bot(query: BotQuery): | |
| # """Main bot endpoint with session management""" | |
| # try: | |
| # # Set global auth context for tools | |
| # global _current_user_id, _current_org_id, _current_auth_token | |
| # _current_user_id = query.userLoginId | |
| # _current_org_id = query.orgId | |
| # _current_auth_token = query.auth_token | |
| # session_id = query.session_id | |
| # # Create new session if not provided | |
| # if not session_id: | |
| # session_data = create_session(query.userLoginId, query.orgId, query.auth_token) | |
| # session_id = session_data["session_id"] | |
| # else: | |
| # # Verify existing session | |
| # get_session(session_id) | |
| # # Add user message to session | |
| # user_message_id = add_message_to_session(session_id, "user", query.message) | |
| # # Create agent with session memory | |
| # agent_executor, memory = create_agent_with_session_memory(session_id) | |
| # # Use the agent to process the query | |
| # result = agent_executor.invoke({"input": query.message}) | |
| # # Add AI response to session | |
| # ai_message_id = add_message_to_session(session_id, "assistant", result["output"]) | |
| # # Update session memory in Redis | |
| # updated_messages = [] | |
| # for message in memory.chat_memory.messages: | |
| # if hasattr(message, 'content'): | |
| # role = "user" if message.__class__.__name__ == "HumanMessage" else "assistant" | |
| # updated_messages.append({ | |
| # "role": role, | |
| # "message": message.content, | |
| # "timestamp": datetime.now().isoformat() | |
| # }) | |
| # update_session_memory(session_id, updated_messages) | |
| # # Clear auth context after use | |
| # _current_user_id = None | |
| # _current_org_id = None | |
| # _current_auth_token = None | |
| # return { | |
| # "session_id": session_id, | |
| # "user_message_id": user_message_id, | |
| # "ai_message_id": ai_message_id, | |
| # "message": query.message, | |
| # "answer": result["output"], | |
| # "userLoginId": query.userLoginId, | |
| # "agent_used": True | |
| # } | |
| # except Exception as e: | |
| # # Clear auth context on error | |
| # _current_user_id = None | |
| # _current_org_id = None | |
| # _current_auth_token = None | |
| # raise HTTPException(status_code=500, detail=f"Error processing chat: {str(e)}") | |
| # @app.get("/sessions/{session_id}/history", response_model=ChatHistoryResponse) | |
| # def get_session_history(session_id: str, n: int = QueryParam(50, description="Number of recent messages to return")): | |
| # """Get chat history for a session""" | |
| # try: | |
| # # Verify session exists | |
| # get_session(session_id) | |
| # # Get chat history | |
| # chat_data = redis_client.get(f"chat:{session_id}") | |
| # if not chat_data: | |
| # return ChatHistoryResponse( | |
| # session_id=session_id, | |
| # messages=[], | |
| # total_messages=0 | |
| # ) | |
| # messages = json.loads(chat_data) | |
| # # Get the last n messages (or all if less than n) | |
| # recent_messages = messages[-n:] if len(messages) > n else messages | |
| # # Convert to MessageResponse objects | |
| # message_responses = [MessageResponse(**msg) for msg in recent_messages] | |
| # return ChatHistoryResponse( | |
| # session_id=session_id, | |
| # messages=message_responses, | |
| # total_messages=len(messages) | |
| # ) | |
| # except Exception as e: | |
| # raise HTTPException(status_code=500, detail=f"Error fetching chat history: {str(e)}") | |
| # @app.post("/chat-documents") | |
| # def chat_documents_only(query: Query): | |
| # """Direct document search without agent""" | |
| # try: | |
| # result = search_documents(query.message) | |
| # return { | |
| # "message": query.message, | |
| # "answer": result, | |
| # "tool_used": "document_search" | |
| # } | |
| # except Exception as e: | |
| # return { | |
| # "message": query.message, | |
| # "answer": f"An error occurred: {str(e)}", | |
| # "tool_used": "document_search" | |
| # } | |
| # @app.post("/list-projects") | |
| # def list_projects(request: ProjectRequest): | |
| # """Direct project listing without agent""" | |
| # try: | |
| # # Use the provided auth token and userLoginId | |
| # encoded_token = get_encoded_auth_token(request.userLoginId, request.auth_token) | |
| # # Fetch projects | |
| # data = fetch_user_projects(request.userLoginId, request.orgId, encoded_token) | |
| # # Format and return the project list | |
| # formatted = format_project_response(data) | |
| # return { | |
| # "projects": formatted, | |
| # "tool_used": "project_list" | |
| # } | |
| # except Exception as e: | |
| # return { | |
| # "error": f"An error occurred: {str(e)}", | |
| # "tool_used": "project_list" | |
| # } | |
| # @app.get("/redis-info") | |
| # def redis_info(): | |
| # """Get Redis connection information""" | |
| # try: | |
| # info = redis_client.info() | |
| # return { | |
| # "redis_connected": True, | |
| # "redis_version": info.get("redis_version"), | |
| # "used_memory": info.get("used_memory_human"), | |
| # "connected_clients": info.get("connected_clients"), | |
| # "total_keys": redis_client.dbsize() | |
| # } | |
| # except Exception as e: | |
| # return { | |
| # "redis_connected": False, | |
| # "error": str(e) | |
| # } | |
| # @app.get("/health") | |
| # def health(): | |
| # try: | |
| # redis_client.ping() | |
| # redis_status = "connected" | |
| # except: | |
| # redis_status = "disconnected" | |
| # return { | |
| # "status": "ok", | |
| # "tools": ["document_search", "project_list"], | |
| # "agent": "active", | |
| # "session_management": "enabled", | |
| # "redis_status": redis_status, | |
| # "total_sessions": len(list(redis_client.scan_iter(match="session:*"))) | |
| # } | |
| # if __name__ == "__main__": | |
| # import uvicorn | |
| # try: | |
| # uvicorn.run(app, host="0.0.0.0", port=8000) | |
| # except KeyboardInterrupt: | |
| # print("\n🛑 Server stopped gracefully") | |
| # except Exception as e: | |
| # print(f"❌ Server error: {e}") | |
| #_______________________________________________session id -- with redis - chat session memory _________________________________________________________ | |
| # | |
| from fastapi import FastAPI, HTTPException, Query as QueryParam | |
| from pydantic import BaseModel, Field | |
| from langchain_openai import ChatOpenAI, OpenAIEmbeddings | |
| from qdrant_client import QdrantClient | |
| from langchain.agents import Tool, AgentExecutor, create_openai_tools_agent | |
| from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder | |
| from langchain.memory import ConversationBufferMemory | |
| from typing import Optional, List, Dict, Any | |
| import os | |
| import warnings | |
| import base64 | |
| import requests | |
| from dotenv import load_dotenv | |
| from datetime import datetime | |
| import json | |
| import uuid | |
| import redis | |
| # Suppress warnings | |
| warnings.filterwarnings("ignore", message="Qdrant client version.*is incompatible.*") | |
| load_dotenv() | |
| app = FastAPI(title="AI Agent with Redis Session Management") | |
| # Environment variables | |
| OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") | |
| QDRANT_COLLECTION_NAME = os.getenv("QDRANT_COLLECTION_NAME", "vatsav_test_1") | |
| QDRANT_HOST = os.getenv("QDRANT_HOST", "127.0.0.1") | |
| QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333)) | |
| # Redis Configuration | |
| REDIS_URL = os.getenv("REDIS_URL") | |
| REDIS_HOST = os.getenv("REDIS_HOST", "127.0.0.1") | |
| REDIS_PORT = int(os.getenv("REDIS_PORT", 6379)) | |
| REDIS_PASSWORD = os.getenv("REDIS_PASSWORD") | |
| # Initialize Redis client | |
| def get_redis_client(): | |
| """Initialize Redis client with fallback to local Redis""" | |
| try: | |
| if REDIS_URL: | |
| # Use deployed Redis URL | |
| redis_client = redis.from_url( | |
| REDIS_URL, | |
| decode_responses=True, | |
| socket_connect_timeout=5, | |
| socket_timeout=5 | |
| ) | |
| # Test connection | |
| redis_client.ping() | |
| print(f"✅ Connected to deployed Redis: {REDIS_URL}") | |
| return redis_client | |
| else: | |
| # Use local Redis | |
| redis_client = redis.StrictRedis( | |
| host=REDIS_HOST, | |
| port=REDIS_PORT, | |
| password=REDIS_PASSWORD, | |
| decode_responses=True, | |
| socket_connect_timeout=5, | |
| socket_timeout=5 | |
| ) | |
| # Test connection | |
| redis_client.ping() | |
| print(f"✅ Connected to local Redis: {REDIS_HOST}:{REDIS_PORT}") | |
| return redis_client | |
| except Exception as e: | |
| print(f"❌ Redis connection failed: {e}") | |
| raise HTTPException(status_code=500, detail=f"Redis connection failed: {str(e)}") | |
| # Initialize Redis client | |
| redis_client = get_redis_client() | |
| # Initialize models | |
| embedding_model = OpenAIEmbeddings( | |
| model="text-embedding-3-large", | |
| openai_api_key=OPENAI_API_KEY, | |
| ) | |
| qdrant_client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT) | |
| llm = ChatOpenAI(model="gpt-4o", temperature=0, openai_api_key=OPENAI_API_KEY) | |
| # === INPUT SCHEMAS === | |
| class Query(BaseModel): | |
| message: str | |
| class ProjectRequest(BaseModel): | |
| userLoginId: int | |
| orgId: int | |
| auth_token: str | |
| class BotQuery(BaseModel): | |
| userLoginId: int | |
| orgId: int | |
| auth_token: str | |
| session_id: Optional[str] = None | |
| message: str | |
| class SessionResponse(BaseModel): | |
| session_id: str | |
| userLoginId: int | |
| orgId: int | |
| created_at: str | |
| status: str | |
| title: Optional[str] = "New Chat" | |
| class MessageResponse(BaseModel): | |
| message_id: str | |
| session_id: str | |
| role: str # "user" or "assistant" | |
| message: str | |
| timestamp: str | |
| class ChatHistoryResponse(BaseModel): | |
| session_id: str | |
| messages: List[MessageResponse] | |
| total_messages: int | |
| # === SESSION MANAGEMENT FUNCTIONS === | |
| def create_session(userLoginId: int, orgId: int, auth_token: str) -> dict: | |
| """Create a new chat session""" | |
| session_id = str(uuid.uuid4()) | |
| session_data = { | |
| "session_id": session_id, | |
| "userLoginId": userLoginId, | |
| "orgId": orgId, | |
| "auth_token": auth_token, | |
| "created_at": datetime.now().isoformat(), | |
| "status": "active", | |
| "title": "New Chat" # Default title, will be updated after first message | |
| } | |
| # Store session in Redis with 24 hour TTL | |
| redis_client.setex( | |
| f"session:{session_id}", | |
| 86400, # 24 hours | |
| json.dumps(session_data) | |
| ) | |
| # Initialize empty chat history | |
| redis_client.setex( | |
| f"chat:{session_id}", | |
| 86400, # 24 hours | |
| json.dumps([]) | |
| ) | |
| # Initialize conversation memory | |
| redis_client.setex( | |
| f"memory:{session_id}", | |
| 86400, # 24 hours | |
| json.dumps([]) | |
| ) | |
| return session_data | |
| def get_session(session_id: str) -> dict: | |
| """Get session data from Redis""" | |
| session_data = redis_client.get(f"session:{session_id}") | |
| if not session_data: | |
| raise HTTPException(status_code=404, detail="Session not found or expired") | |
| return json.loads(session_data) | |
| def add_message_to_session(session_id: str, role: str, message: str) -> str: | |
| """Add message to session chat history""" | |
| message_id = str(uuid.uuid4()) | |
| message_data = { | |
| "message_id": message_id, | |
| "session_id": session_id, | |
| "role": role, | |
| "message": message, | |
| "timestamp": datetime.now().isoformat() | |
| } | |
| # Get current chat history | |
| chat_history = redis_client.get(f"chat:{session_id}") | |
| if chat_history: | |
| messages = json.loads(chat_history) | |
| else: | |
| messages = [] | |
| # Add new message | |
| messages.append(message_data) | |
| # Update chat history in Redis with extended TTL | |
| redis_client.setex( | |
| f"chat:{session_id}", | |
| 86400, # 24 hours | |
| json.dumps(messages) | |
| ) | |
| return message_id | |
| def get_session_memory(session_id: str) -> List[Dict]: | |
| """Get conversation memory for session""" | |
| memory_data = redis_client.get(f"memory:{session_id}") | |
| if memory_data: | |
| return json.loads(memory_data) | |
| return [] | |
| def update_session_memory(session_id: str, messages: List[Dict]): | |
| """Update conversation memory for session""" | |
| redis_client.setex( | |
| f"memory:{session_id}", | |
| 86400, # 24 hours | |
| json.dumps(messages) | |
| ) | |
| def update_session_title(session_id: str): | |
| """Update session title after first message""" | |
| try: | |
| # Get session data | |
| session_data = redis_client.get(f"session:{session_id}") | |
| if not session_data: | |
| return | |
| session = json.loads(session_data) | |
| # Only update if current title is "New Chat" | |
| if session.get("title", "New Chat") == "New Chat": | |
| new_title = generate_session_title(session_id) | |
| session["title"] = new_title | |
| # Update session in Redis | |
| redis_client.setex( | |
| f"session:{session_id}", | |
| 86400, # 24 hours | |
| json.dumps(session) | |
| ) | |
| except Exception as e: | |
| print(f"Error updating session title: {e}") | |
| pass # Don't fail the request if title update fails | |
| def generate_session_title(session_id: str) -> str: | |
| """Generate a title for the session based on chat history""" | |
| try: | |
| # Get chat history | |
| chat_data = redis_client.get(f"chat:{session_id}") | |
| if not chat_data: | |
| return "New Chat" | |
| messages = json.loads(chat_data) | |
| if not messages: | |
| return "New Chat" | |
| # Get first user message for title generation | |
| first_user_message = None | |
| for msg in messages: | |
| if msg["role"] == "user": | |
| first_user_message = msg["message"] | |
| break | |
| if not first_user_message: | |
| return "New Chat" | |
| # Generate title using LLM | |
| title_prompt = f"""Generate a short, descriptive title (maximum 6 words) for a chat conversation that starts with this message: | |
| "{first_user_message[:200]}" | |
| Return only the title, no quotes or additional text. The title should capture the main topic or intent of the conversation.""" | |
| try: | |
| response = llm.invoke(title_prompt) | |
| title = response.content.strip() | |
| # Clean and limit title | |
| title = title.replace('"', '').replace("'", "") | |
| if len(title) > 50: | |
| title = title[:47] + "..." | |
| return title if title else "New Chat" | |
| except Exception as e: | |
| print(f"Error generating title: {e}") | |
| # Fallback: use first few words of the message | |
| words = first_user_message.split()[:4] | |
| return " ".join(words) + ("..." if len(words) >= 4 else "") | |
| except Exception as e: | |
| print(f"Error in generate_session_title: {e}") | |
| return "New Chat" | |
| def get_user_sessions(userLoginId: int) -> List[dict]: | |
| """Get all sessions for a user with generated titles""" | |
| sessions = [] | |
| # Scan for all session keys | |
| for key in redis_client.scan_iter(match="session:*"): | |
| session_data = redis_client.get(key) | |
| if session_data: | |
| session = json.loads(session_data) | |
| if session["userLoginId"] == userLoginId: | |
| # Generate title based on chat history | |
| session["title"] = generate_session_title(session["session_id"]) | |
| sessions.append(session) | |
| # Sort sessions by created_at (most recent first) | |
| sessions.sort(key=lambda x: x["created_at"], reverse=True) | |
| return sessions | |
| def delete_session(session_id: str): | |
| """Delete session and associated data""" | |
| # Delete session data | |
| redis_client.delete(f"session:{session_id}") | |
| # Delete chat history | |
| redis_client.delete(f"chat:{session_id}") | |
| # Delete memory | |
| redis_client.delete(f"memory:{session_id}") | |
| # === UTILITY FUNCTIONS === | |
| def get_encoded_auth_token(user: int, token: str) -> str: | |
| auth_string = f"{user}:{token}" | |
| return base64.b64encode(auth_string.encode("utf-8")).decode("utf-8") | |
| def fetch_user_projects(userLoginId: int, orgId: int, auth_token: str): | |
| url = "https://japidemo.dev.ingenspark.com/fetchUserProjects" | |
| payload = { | |
| "userLoginId": userLoginId, | |
| "orgId": orgId | |
| } | |
| headers = { | |
| 'accept': 'application/json, text/plain, */*', | |
| 'authorization': f'Basic {auth_token}', | |
| 'content-type': 'application/json; charset=UTF-8' | |
| } | |
| try: | |
| response = requests.post(url, headers=headers, json=payload) | |
| response.raise_for_status() | |
| return response.json() | |
| except requests.exceptions.RequestException as e: | |
| raise HTTPException(status_code=response.status_code if 'response' in locals() else 500, | |
| detail=str(e)) | |
| def format_project_response(data: dict) -> str: | |
| my_projects = data.get("data", {}).get("Myprojects", []) | |
| other_projects = data.get("data", {}).get("Otherprojects", []) | |
| all_projects = [] | |
| for project in my_projects: | |
| all_projects.append({ | |
| "type": "Your Project", | |
| "projectNm": project["projectNm"], | |
| "projectId": project["projectId"], | |
| "created_dttm": project["created_dttm"].split('.')[0], | |
| "description": project["description"], | |
| "categoryName": project["categoryName"] | |
| }) | |
| for project in other_projects: | |
| all_projects.append({ | |
| "type": "Other Project", | |
| "projectNm": project["projectNm"], | |
| "projectId": project["projectId"], | |
| "created_dttm": project["created_dttm"].split('.')[0], | |
| "description": project["description"], | |
| "categoryName": project["categoryName"] | |
| }) | |
| if not all_projects: | |
| return "❌ No projects found." | |
| # Build the formatted string | |
| result = [f"✅ You have access to {len(all_projects)} project(s):\n"] | |
| for i, project in enumerate(all_projects, 1): | |
| result.append(f"{i}. Project Name: {project['projectNm']} ({project['type']})") | |
| result.append(f" Project ID: {project['projectId']}") | |
| result.append(f" Created On: {project['created_dttm']}") | |
| result.append(f" Description: {project['description']}") | |
| result.append(f" Category: {project['categoryName']}\n") | |
| return "\n".join(result) | |
| # === TOOL FUNCTIONS === | |
| def search_documents(query: str) -> str: | |
| """Search through ingested documents and get relevant information.""" | |
| try: | |
| # Generate embedding for the query | |
| query_vector = embedding_model.embed_query(query) | |
| # Search in Qdrant | |
| search_result = qdrant_client.search( | |
| collection_name=QDRANT_COLLECTION_NAME, | |
| query_vector=query_vector, | |
| limit=5, | |
| ) | |
| if not search_result: | |
| return "No relevant information found in the knowledge base." | |
| # Convert results to text content | |
| context_texts = [] | |
| sources = [] | |
| for hit in search_result: | |
| context_texts.append(hit.payload["text"]) | |
| sources.append(hit.payload.get("source", "unknown")) | |
| # Create a simple prompt for answering based on context | |
| context = "\n\n".join(context_texts) | |
| unique_sources = list(set(sources)) | |
| # Use the LLM directly to answer the message based on context | |
| prompt = f"""Based on the following context, answer the message: {query} | |
| Context: | |
| {context} | |
| Please provide a comprehensive answer based on the context above. If the context doesn't contain enough information to answer the message, say so clearly.""" | |
| response = llm.invoke(prompt) | |
| return f"{response.content}\n\nSources: {', '.join(unique_sources)}" | |
| except Exception as e: | |
| return f"Error searching documents: {str(e)}" | |
| # Global variables to store auth context (for tool functions) | |
| _current_user_id = None | |
| _current_org_id = None | |
| _current_auth_token = None | |
| def get_user_projects(userLoginId: str) -> str: | |
| """Get list of projects for a user.""" | |
| try: | |
| # Use global auth context if available | |
| if _current_auth_token and _current_user_id: | |
| user_id = _current_user_id | |
| org_id = _current_org_id or 1 | |
| auth_token = _current_auth_token | |
| else: | |
| return "❌ Authentication token required. Please provide auth_token in your request." | |
| # Encode auth token using the actual user ID and provided token | |
| encoded_token = get_encoded_auth_token(user_id, auth_token) | |
| # Fetch projects | |
| data = fetch_user_projects(user_id, org_id, encoded_token) | |
| # Format and return the project list | |
| formatted = format_project_response(data) | |
| return formatted | |
| except ValueError: | |
| return "❌ Invalid userLoginId format. Please provide a valid number." | |
| except Exception as e: | |
| return f"❌ Error fetching projects: {str(e)}" | |
| # === CREATE TOOLS === | |
| document_search_tool = Tool( | |
| name="document_search", | |
| description="""Use this tool to search through ingested documents and get relevant information from the knowledge base. | |
| Perfect for answering messages about uploaded documents, manuals, or any content that was previously stored. | |
| Input should be a search query or message about the documents.""", | |
| func=search_documents | |
| ) | |
| project_list_tool = Tool( | |
| name="get_user_projects", | |
| description="""Use this tool to get the list of projects for a user. | |
| Perfect for when users ask about their projects, want to see available projects, or need project information. | |
| Input should be the userLoginId (e.g., '25'). | |
| Note: This tool requires authentication context to be set.""", | |
| func=get_user_projects | |
| ) | |
| # === AGENT SETUP === | |
| def create_agent_with_session_memory(session_id: str): | |
| """Create agent with session memory from Redis""" | |
| # Get memory from Redis | |
| memory_messages = get_session_memory(session_id) | |
| agent_prompt = ChatPromptTemplate.from_messages([ | |
| ("system", """You are a helpful AI assistant with access to multiple tools and conversation memory: | |
| 1. **Document Search**: Search through uploaded documents and knowledge base | |
| 2. **Project Management**: Get list of user projects and project information | |
| Your capabilities: | |
| - Answer messages about documents using the document search tool | |
| - Help users find their projects and project information | |
| - Remember previous conversations in this session | |
| - Provide general assistance and information | |
| - Use appropriate tools based on user queries | |
| Guidelines: | |
| - Use the document search tool when users ask about specific content, documentation, or information that might be in uploaded files | |
| - Use the project tool when users ask about projects, want to see their projects, or need project-related information | |
| - Reference previous conversation context when relevant | |
| - Be clear about which tool you're using and what information you're providing | |
| - If you're unsure which tool to use, you can ask for clarification | |
| - Provide helpful, accurate, and well-formatted responses | |
| Remember: Always use the most appropriate tool based on the user's message and conversation context to provide the best possible answer."""), | |
| MessagesPlaceholder(variable_name="chat_history"), | |
| ("user", "{input}"), | |
| MessagesPlaceholder(variable_name="agent_scratchpad"), | |
| ]) | |
| # Create memory object | |
| memory = ConversationBufferMemory( | |
| memory_key="chat_history", | |
| return_messages=True | |
| ) | |
| # Load existing messages into memory | |
| for msg in memory_messages: | |
| if msg["role"] == "user": | |
| memory.chat_memory.add_user_message(msg["message"]) | |
| else: | |
| memory.chat_memory.add_ai_message(msg["message"]) | |
| # Create tools list | |
| tools = [document_search_tool, project_list_tool] | |
| # Create the agent | |
| agent = create_openai_tools_agent(llm, tools, agent_prompt) | |
| # Create the agent executor with memory | |
| agent_executor = AgentExecutor( | |
| agent=agent, | |
| tools=tools, | |
| verbose=True, | |
| memory=memory | |
| ) | |
| return agent_executor, memory | |
| # === API ENDPOINTS === | |
| def create_new_session(userLoginId: int, orgId: int, auth_token: str): | |
| """Create a new chat session""" | |
| try: | |
| session_data = create_session(userLoginId, orgId, auth_token) | |
| return SessionResponse(**session_data) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Error creating session: {str(e)}") | |
| def list_user_sessions(userLoginId: int): | |
| """List all sessions for a user""" | |
| try: | |
| sessions = get_user_sessions(userLoginId) | |
| return { | |
| "userLoginId": userLoginId, | |
| "total_sessions": len(sessions), | |
| "sessions": sessions | |
| } | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Error fetching sessions: {str(e)}") | |
| def delete_user_session(session_id: str): | |
| """Delete/close a session""" | |
| try: | |
| # Verify session exists | |
| get_session(session_id) | |
| # Delete session | |
| delete_session(session_id) | |
| return { | |
| "message": f"Session {session_id} deleted successfully", | |
| "session_id": session_id | |
| } | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Error deleting session: {str(e)}") | |
| def chat_with_bot(query: BotQuery): | |
| """Main bot endpoint with session management""" | |
| try: | |
| # Set global auth context for tools | |
| global _current_user_id, _current_org_id, _current_auth_token | |
| _current_user_id = query.userLoginId | |
| _current_org_id = query.orgId | |
| _current_auth_token = query.auth_token | |
| session_id = query.session_id | |
| # Create new session if not provided | |
| if not session_id: | |
| session_data = create_session(query.userLoginId, query.orgId, query.auth_token) | |
| session_id = session_data["session_id"] | |
| else: | |
| # Verify existing session | |
| get_session(session_id) | |
| # Add user message to session | |
| user_message_id = add_message_to_session(session_id, "user", query.message) | |
| # Create agent with session memory | |
| agent_executor, memory = create_agent_with_session_memory(session_id) | |
| # Use the agent to process the query | |
| result = agent_executor.invoke({"input": query.message}) | |
| # Add AI response to session | |
| ai_message_id = add_message_to_session(session_id, "assistant", result["output"]) | |
| # Update session memory in Redis | |
| updated_messages = [] | |
| for message in memory.chat_memory.messages: | |
| if hasattr(message, 'content'): | |
| role = "user" if message.__class__.__name__ == "HumanMessage" else "assistant" | |
| updated_messages.append({ | |
| "role": role, | |
| "message": message.content, | |
| "timestamp": datetime.now().isoformat() | |
| }) | |
| update_session_memory(session_id, updated_messages) | |
| # Update session title after first user message | |
| update_session_title(session_id) | |
| # Clear auth context after use | |
| _current_user_id = None | |
| _current_org_id = None | |
| _current_auth_token = None | |
| return { | |
| "session_id": session_id, | |
| "user_message_id": user_message_id, | |
| "ai_message_id": ai_message_id, | |
| "message": query.message, | |
| "answer": result["output"], | |
| "userLoginId": query.userLoginId, | |
| "agent_used": True | |
| } | |
| except Exception as e: | |
| # Clear auth context on error | |
| _current_user_id = None | |
| _current_org_id = None | |
| _current_auth_token = None | |
| raise HTTPException(status_code=500, detail=f"Error processing chat: {str(e)}") | |
| def get_session_history(session_id: str, n: int = QueryParam(50, description="Number of recent messages to return")): | |
| """Get chat history for a session""" | |
| try: | |
| # Verify session exists | |
| get_session(session_id) | |
| # Get chat history | |
| chat_data = redis_client.get(f"chat:{session_id}") | |
| if not chat_data: | |
| return ChatHistoryResponse( | |
| session_id=session_id, | |
| messages=[], | |
| total_messages=0 | |
| ) | |
| messages = json.loads(chat_data) | |
| # Get the last n messages (or all if less than n) | |
| recent_messages = messages[-n:] if len(messages) > n else messages | |
| # Convert to MessageResponse objects | |
| message_responses = [MessageResponse(**msg) for msg in recent_messages] | |
| return ChatHistoryResponse( | |
| session_id=session_id, | |
| messages=message_responses, | |
| total_messages=len(messages) | |
| ) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Error fetching chat history: {str(e)}") | |
| def chat_documents_only(query: Query): | |
| """Direct document search without agent""" | |
| try: | |
| result = search_documents(query.message) | |
| return { | |
| "message": query.message, | |
| "answer": result, | |
| "tool_used": "document_search" | |
| } | |
| except Exception as e: | |
| return { | |
| "message": query.message, | |
| "answer": f"An error occurred: {str(e)}", | |
| "tool_used": "document_search" | |
| } | |
| def list_projects(request: ProjectRequest): | |
| """Direct project listing without agent""" | |
| try: | |
| # Use the provided auth token and userLoginId | |
| encoded_token = get_encoded_auth_token(request.userLoginId, request.auth_token) | |
| # Fetch projects | |
| data = fetch_user_projects(request.userLoginId, request.orgId, encoded_token) | |
| # Format and return the project list | |
| formatted = format_project_response(data) | |
| return { | |
| "projects": formatted, | |
| "tool_used": "project_list" | |
| } | |
| except Exception as e: | |
| return { | |
| "error": f"An error occurred: {str(e)}", | |
| "tool_used": "project_list" | |
| } | |
| def refresh_session_title(session_id: str): | |
| """Manually refresh/regenerate session title""" | |
| try: | |
| # Verify session exists | |
| session_data = get_session(session_id) | |
| # Generate new title | |
| new_title = generate_session_title(session_id) | |
| # Update session | |
| session_data["title"] = new_title | |
| redis_client.setex( | |
| f"session:{session_id}", | |
| 86400, # 24 hours | |
| json.dumps(session_data) | |
| ) | |
| return { | |
| "session_id": session_id, | |
| "new_title": new_title, | |
| "message": "Session title updated successfully" | |
| } | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Error updating session title: {str(e)}") | |
| def redis_info(): | |
| """Get Redis connection information""" | |
| try: | |
| info = redis_client.info() | |
| return { | |
| "redis_connected": True, | |
| "redis_version": info.get("redis_version"), | |
| "used_memory": info.get("used_memory_human"), | |
| "connected_clients": info.get("connected_clients"), | |
| "total_keys": redis_client.dbsize() | |
| } | |
| except Exception as e: | |
| return { | |
| "redis_connected": False, | |
| "error": str(e) | |
| } | |
| def health(): | |
| try: | |
| redis_client.ping() | |
| redis_status = "connected" | |
| except: | |
| redis_status = "disconnected" | |
| return { | |
| "status": "ok", | |
| "tools": ["document_search", "project_list"], | |
| "agent": "active", | |
| "session_management": "enabled", | |
| "redis_status": redis_status, | |
| "total_sessions": len(list(redis_client.scan_iter(match="session:*"))) | |
| } | |
| if __name__ == "__main__": | |
| import uvicorn | |
| try: | |
| uvicorn.run(app, host="0.0.0.0", port=8000) | |
| except KeyboardInterrupt: | |
| print("\n🛑 Server stopped gracefully") | |
| except Exception as e: | |
| print(f"❌ Server error: {e}") | |
| #______________________________________________sucess true false _pandasai tool_________________________________________________________________ | |
| # from dotenv import load_dotenv | |
| # from datetime import datetime | |
| # import json | |
| # import uuid | |
| # import redis | |
| # # Pandas AI imports | |
| # import re | |
| # import urllib.parse | |
| # import pandas as pd | |
| # import dask.dataframe as dd | |
| # from math import ceil | |
| # import psycopg2 | |
| # from pandasai import SmartDataframe | |
| # from pandasai.llm.openai import OpenAI as PandasOpenAI | |
| # from fastapi import FastAPI, Request | |
| # from fastapi.responses import JSONResponse | |
| # import json | |
| # # Import your existing S3 connection details | |
| # from retrive_secrects import * # CONNECTIONS_HOST, etc. | |
| # # Suppress warnings | |
| # warnings.filterwarnings("ignore", message="Qdrant client version.*is incompatible.*") | |
| # load_dotenv() | |
| # app = FastAPI(title="AI Agent with Redis Session Management and Pandas AI") | |
| # # Environment variables | |
| # OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") | |
| # QDRANT_COLLECTION_NAME = os.getenv("QDRANT_COLLECTION_NAME", "vatsav_test_1") | |
| # QDRANT_HOST = os.getenv("QDRANT_HOST", "127.0.0.1") | |
| # QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333)) | |
| # # Redis Configuration | |
| # REDIS_URL = os.getenv("REDIS_URL") | |
| # REDIS_HOST = os.getenv("REDIS_HOST", "127.0.0.1") | |
| # REDIS_PORT = int(os.getenv("REDIS_PORT", 6379)) | |
| # REDIS_PASSWORD = os.getenv("REDIS_PASSWORD") | |
| # # S3 Constants (from your original code) | |
| # S3_Bucket_Name = 'ingenspark-user-files' | |
| # S3_Raw_Files_Folder = 'User-Uploaded-Raw-Files' | |
| # S3_Modified_Files_Folder = 'Modified-Files/' | |
| # S3_Output_Files_Folder = 'Output-Files/' | |
| # S3_Published_Results_Folder = 'Published-Results/' | |
| # S3_Ingen_Customer_Output = 'Ingen-Customer/' | |
| # Dominant_Segmentation_Output = 'Dominant-Segmentation/' | |
| # Trend_Segmentation_Output = 'Trend-Segmentation/' | |
| # Decile_Quartile_segmentation_Output = 'Decile-Quartile-Segmentation/' | |
| # Combined_Segmentation_Output = 'Combine-Segmentation/' | |
| # Custom_Segmentation_Output = 'Custom-Segmentation/' | |
| # Customer_360_Output = 'Customer-360/' | |
| # Merge_file_folder = S3_Modified_Files_Folder + 'IngenData-Merged-Tables/' | |
| # S3_Dev_Doc_Images_Folder = 'Developers-Documentation-Images/' | |
| # S3_Temporary_Files_Folder = S3_Raw_Files_Folder | |
| # S3_App_Specific_Data = 'Application-Specific-Data/' | |
| # S3_Transformation_Tables_Folder = 'Modified-Files/Modified-Tables/Transformation-Tables/' | |
| # cloud_front_url = "https://files.dev.ingenspark.com/" | |
| # # Initialize Redis client | |
| # def get_redis_client(): | |
| # """Initialize Redis client with fallback to local Redis""" | |
| # try: | |
| # if REDIS_URL: | |
| # # Use deployed Redis URL | |
| # redis_client = redis.from_url( | |
| # REDIS_URL, | |
| # decode_responses=True, | |
| # socket_connect_timeout=5, | |
| # socket_timeout=5 | |
| # ) | |
| # # Test connection | |
| # redis_client.ping() | |
| # print(f"✅ Connected to deployed Redis: {REDIS_URL}") | |
| # return redis_client | |
| # else: | |
| # # Use local Redis | |
| # redis_client = redis.StrictRedis( | |
| # host=REDIS_HOST, | |
| # port=REDIS_PORT, | |
| # password=REDIS_PASSWORD, | |
| # decode_responses=True, | |
| # socket_connect_timeout=5, | |
| # socket_timeout=5 | |
| # ) | |
| # # Test connection | |
| # redis_client.ping() | |
| # print(f"✅ Connected to local Redis: {REDIS_HOST}:{REDIS_PORT}") | |
| # return redis_client | |
| # except Exception as e: | |
| # print(f"❌ Redis connection failed: {e}") | |
| # raise HTTPException(status_code=500, detail=f"Redis connection failed: {str(e)}") | |
| # # Initialize Redis client | |
| # redis_client = get_redis_client() | |
| # # Initialize models | |
| # embedding_model = OpenAIEmbeddings( | |
| # model="text-embedding-3-large", | |
| # openai_api_key=OPENAI_API_KEY, | |
| # ) | |
| # qdrant_client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT) | |
| # llm = ChatOpenAI(model="gpt-4o", temperature=0, openai_api_key=OPENAI_API_KEY) | |
| # # === PANDAS AI FUNCTIONS === | |
| # def read_parquet_file_from_s3(ufuid=None, columns_list=None, records_count=None, file_location=''): | |
| # """ | |
| # Reads a Parquet file from S3 using Dask and returns it as a Pandas DataFrame. | |
| # Parameters: | |
| # ufuid (int): Optional user_file_upload_id to fetch S3 path from DB. | |
| # columns_list (list/str): Columns to read. | |
| # records_count (int): Not used currently. | |
| # file_location (str): Direct file path in S3. | |
| # Returns: | |
| # pandas.DataFrame | |
| # """ | |
| # try: | |
| # # Connect to PostgreSQL | |
| # conn = psycopg2.connect( | |
| # host=CONNECTIONS_HOST, | |
| # database=CONNECTIONS_DB, | |
| # user=CONNECTIONS_USER, | |
| # password=CONNECTIONS_PASS | |
| # ) | |
| # cursor = conn.cursor() | |
| # if ufuid is not None: | |
| # query = """SELECT file_name, table_names FROM public.user_file_upload WHERE user_file_upload_id = %s""" | |
| # cursor.execute(query, (ufuid,)) | |
| # file = cursor.fetchone() | |
| # if not file: | |
| # raise ValueError(f"No file found for ufuid: {ufuid}") | |
| # file_name, s3_file_path = file | |
| # else: | |
| # # Normalize input path | |
| # file_location = re.sub(r'\.parquet(?!$)', '', file_location) | |
| # s3_file_path = file_location if file_location.endswith('.parquet') else file_location + '.parquet' | |
| # # Extract relative S3 path | |
| # s3_file_path = urllib.parse.unquote(s3_file_path.split(f"{S3_Bucket_Name}/")[-1]) | |
| # if not s3_file_path.endswith('.parquet'): | |
| # s3_file_path += '.parquet' | |
| # # Parse columns if given as comma-separated string | |
| # if columns_list and not isinstance(columns_list, list): | |
| # columns_list = [col.strip(' "\'') for col in columns_list.split(',')] | |
| # print(f"\n{'!' * 100}\nReading from: s3://{S3_Bucket_Name}/{s3_file_path}\n") | |
| # # Read using Dask | |
| # ddf = dd.read_parquet( | |
| # f"s3://{S3_Bucket_Name}/{s3_file_path}", | |
| # engine="pyarrow", | |
| # columns=columns_list, | |
| # assume_missing=True | |
| # ) | |
| # ddf = ddf.repartition(npartitions=8) # Optimize for processing | |
| # print("Reading Parquet file from S3 completed successfully.") | |
| # # Close database connection | |
| # cursor.close() | |
| # conn.close() | |
| # return ddf.compute() | |
| # except Exception as e: | |
| # print(f"❌ Error reading Parquet file: {e}") | |
| # return pd.DataFrame() # Return empty DataFrame on error | |
| # def pandas_agent(filepath: str, query: str) -> str: | |
| # """ | |
| # PandasAI agent that reads data from S3 and answers queries about the data. | |
| # Parameters: | |
| # filepath (str): S3 file path or ufuid | |
| # query (str): Natural language query about the data | |
| # Returns: | |
| # str: Answer from PandasAI | |
| # """ | |
| # try: | |
| # # Check if filepath is a number (ufuid) or a file path | |
| # if filepath.isdigit(): | |
| # # It's a ufuid | |
| # data = read_parquet_file_from_s3(ufuid=int(filepath)) | |
| # else: | |
| # # It's a file path | |
| # data = read_parquet_file_from_s3(file_location=filepath) | |
| # if data.empty: | |
| # return "❌ No data found or failed to load the file. Please check the file path or ufuid." | |
| # # Initialize PandasAI LLM | |
| # if not OPENAI_API_KEY: | |
| # return "❌ OPENAI_API_KEY is not set in environment variables." | |
| # pandas_llm = PandasOpenAI(api_token=OPENAI_API_KEY) | |
| # # Create SmartDataframe | |
| # sdf = SmartDataframe(data, config={"llm": pandas_llm}) | |
| # # Ask the question | |
| # print(f"🔍 Processing query: {query}") | |
| # result = sdf.chat(query) | |
| # # Handle different types of results | |
| # if isinstance(result, str): | |
| # return f"📊 Analysis Result:\n{result}" | |
| # elif isinstance(result, (pd.DataFrame, pd.Series)): | |
| # return f"📊 Analysis Result:\n{result.to_string()}" | |
| # else: | |
| # return f"📊 Analysis Result:\n{str(result)}" | |
| # except Exception as e: | |
| # error_msg = f"❌ Error in pandas_agent: {str(e)}" | |
| # print(error_msg) | |
| # return error_msg | |
| # # === INPUT SCHEMAS === | |
| # class Query(BaseModel): | |
| # message: str | |
| # class ProjectRequest(BaseModel): | |
| # userLoginId: int | |
| # orgId: int | |
| # auth_token: str | |
| # class BotQuery(BaseModel): | |
| # userLoginId: int | |
| # orgId: int | |
| # auth_token: str | |
| # session_id: Optional[str] = None | |
| # message: str | |
| # class PandasAgentQuery(BaseModel): | |
| # filepath: str = Field(..., description="S3 file path or ufuid") | |
| # query: str = Field(..., description="Natural language query about the data") | |
| # class SessionResponse(BaseModel): | |
| # session_id: str | |
| # userLoginId: int | |
| # orgId: int | |
| # created_at: str | |
| # status: str | |
| # title: Optional[str] = "New Chat" | |
| # class MessageResponse(BaseModel): | |
| # message_id: str | |
| # session_id: str | |
| # role: str # "user" or "assistant" | |
| # message: str | |
| # timestamp: str | |
| # class ChatHistoryResponse(BaseModel): | |
| # session_id: str | |
| # messages: List[MessageResponse] | |
| # total_messages: int | |
| # # === SESSION MANAGEMENT FUNCTIONS === | |
| # @app.middleware("http") | |
| # async def add_success_flag(request: Request, call_next): | |
| # response = await call_next(request) | |
| # # Only modify JSON responses | |
| # if "application/json" in response.headers.get("content-type", ""): | |
| # try: | |
| # body = b"".join([chunk async for chunk in response.body_iterator]) | |
| # data = json.loads(body.decode()) | |
| # # Add success flag | |
| # data["success"] = 200 <= response.status_code < 300 | |
| # # Build new JSONResponse (auto handles Content-Length) | |
| # response = JSONResponse( | |
| # content=data, | |
| # status_code=response.status_code, | |
| # headers={k: v for k, v in response.headers.items() if k.lower() != "content-length"}, | |
| # ) | |
| # except Exception: | |
| # # fallback if response is not JSON parseable | |
| # pass | |
| # return response | |
| # def create_session(userLoginId: int, orgId: int, auth_token: str) -> dict: | |
| # """Create a new chat session""" | |
| # session_id = str(uuid.uuid4()) | |
| # session_data = { | |
| # "session_id": session_id, | |
| # "userLoginId": userLoginId, | |
| # "orgId": orgId, | |
| # "auth_token": auth_token, | |
| # "created_at": datetime.now().isoformat(), | |
| # "status": "active", | |
| # "title": "New Chat" # Default title, will be updated after first message | |
| # } | |
| # # Store session in Redis with 24 hour TTL | |
| # redis_client.setex( | |
| # f"session:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps(session_data) | |
| # ) | |
| # # Initialize empty chat history | |
| # redis_client.setex( | |
| # f"chat:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps([]) | |
| # ) | |
| # # Initialize conversation memory | |
| # redis_client.setex( | |
| # f"memory:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps([]) | |
| # ) | |
| # return session_data | |
| # def get_session(session_id: str) -> dict: | |
| # """Get session data from Redis""" | |
| # session_data = redis_client.get(f"session:{session_id}") | |
| # if not session_data: | |
| # raise HTTPException(status_code=404, detail="Session not found or expired") | |
| # return json.loads(session_data) | |
| # def add_message_to_session(session_id: str, role: str, message: str) -> str: | |
| # """Add message to session chat history""" | |
| # message_id = str(uuid.uuid4()) | |
| # message_data = { | |
| # "message_id": message_id, | |
| # "session_id": session_id, | |
| # "role": role, | |
| # "message": message, | |
| # "timestamp": datetime.now().isoformat() | |
| # } | |
| # # Get current chat history | |
| # chat_history = redis_client.get(f"chat:{session_id}") | |
| # if chat_history: | |
| # messages = json.loads(chat_history) | |
| # else: | |
| # messages = [] | |
| # # Add new message | |
| # messages.append(message_data) | |
| # # Update chat history in Redis with extended TTL | |
| # redis_client.setex( | |
| # f"chat:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps(messages) | |
| # ) | |
| # return message_id | |
| # def get_session_memory(session_id: str) -> List[Dict]: | |
| # """Get conversation memory for session""" | |
| # memory_data = redis_client.get(f"memory:{session_id}") | |
| # if memory_data: | |
| # return json.loads(memory_data) | |
| # return [] | |
| # def update_session_memory(session_id: str, messages: List[Dict]): | |
| # """Update conversation memory for session""" | |
| # redis_client.setex( | |
| # f"memory:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps(messages) | |
| # ) | |
| # def update_session_title(session_id: str): | |
| # """Update session title after first message""" | |
| # try: | |
| # # Get session data | |
| # session_data = redis_client.get(f"session:{session_id}") | |
| # if not session_data: | |
| # return | |
| # session = json.loads(session_data) | |
| # # Only update if current title is "New Chat" | |
| # if session.get("title", "New Chat") == "New Chat": | |
| # new_title = generate_session_title(session_id) | |
| # session["title"] = new_title | |
| # # Update session in Redis | |
| # redis_client.setex( | |
| # f"session:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps(session) | |
| # ) | |
| # except Exception as e: | |
| # print(f"Error updating session title: {e}") | |
| # pass # Don't fail the request if title update fails | |
| # def generate_session_title(session_id: str) -> str: | |
| # """Generate a title for the session based on chat history""" | |
| # try: | |
| # # Get chat history | |
| # chat_data = redis_client.get(f"chat:{session_id}") | |
| # if not chat_data: | |
| # return "New Chat" | |
| # messages = json.loads(chat_data) | |
| # if not messages: | |
| # return "New Chat" | |
| # # Get first user message for title generation | |
| # first_user_message = None | |
| # for msg in messages: | |
| # if msg["role"] == "user": | |
| # first_user_message = msg["message"] | |
| # break | |
| # if not first_user_message: | |
| # return "New Chat" | |
| # # Generate title using LLM | |
| # title_prompt = f"""Generate a short, descriptive title (maximum 6 words) for a chat conversation that starts with this message: | |
| # "{first_user_message[:200]}" | |
| # Return only the title, no quotes or additional text. The title should capture the main topic or intent of the conversation.""" | |
| # try: | |
| # response = llm.invoke(title_prompt) | |
| # title = response.content.strip() | |
| # # Clean and limit title | |
| # title = title.replace('"', '').replace("'", "") | |
| # if len(title) > 50: | |
| # title = title[:47] + "..." | |
| # return title if title else "New Chat" | |
| # except Exception as e: | |
| # print(f"Error generating title: {e}") | |
| # # Fallback: use first few words of the message | |
| # words = first_user_message.split()[:4] | |
| # return " ".join(words) + ("..." if len(words) >= 4 else "") | |
| # except Exception as e: | |
| # print(f"Error in generate_session_title: {e}") | |
| # return "New Chat" | |
| # def get_user_sessions(userLoginId: int) -> List[dict]: | |
| # """Get all sessions for a user with generated titles""" | |
| # sessions = [] | |
| # # Scan for all session keys | |
| # for key in redis_client.scan_iter(match="session:*"): | |
| # session_data = redis_client.get(key) | |
| # if session_data: | |
| # session = json.loads(session_data) | |
| # if session["userLoginId"] == userLoginId: | |
| # # Generate title based on chat history | |
| # session["title"] = generate_session_title(session["session_id"]) | |
| # sessions.append(session) | |
| # # Sort sessions by created_at (most recent first) | |
| # sessions.sort(key=lambda x: x["created_at"], reverse=True) | |
| # return sessions | |
| # def delete_session(session_id: str): | |
| # """Delete session and associated data""" | |
| # # Delete session data | |
| # redis_client.delete(f"session:{session_id}") | |
| # # Delete chat history | |
| # redis_client.delete(f"chat:{session_id}") | |
| # # Delete memory | |
| # redis_client.delete(f"memory:{session_id}") | |
| # # === UTILITY FUNCTIONS === | |
| # def get_encoded_auth_token(user: int, token: str) -> str: | |
| # auth_string = f"{user}:{token}" | |
| # return base64.b64encode(auth_string.encode("utf-8")).decode("utf-8") | |
| # def fetch_user_projects(userLoginId: int, orgId: int, auth_token: str): | |
| # url = "https://japidemo.dev.ingenspark.com/fetchUserProjects" | |
| # payload = { | |
| # "userLoginId": userLoginId, | |
| # "orgId": orgId | |
| # } | |
| # headers = { | |
| # 'accept': 'application/json, text/plain, */*', | |
| # 'authorization': f'Basic {auth_token}', | |
| # 'content-type': 'application/json; charset=UTF-8' | |
| # } | |
| # try: | |
| # response = requests.post(url, headers=headers, json=payload) | |
| # response.raise_for_status() | |
| # return response.json() | |
| # except requests.exceptions.RequestException as e: | |
| # raise HTTPException(status_code=response.status_code if 'response' in locals() else 500, | |
| # detail=str(e)) | |
| # def format_project_response(data: dict) -> str: | |
| # my_projects = data.get("data", {}).get("Myprojects", []) | |
| # other_projects = data.get("data", {}).get("Otherprojects", []) | |
| # all_projects = [] | |
| # for project in my_projects: | |
| # all_projects.append({ | |
| # "type": "Your Project", | |
| # "projectNm": project["projectNm"], | |
| # "projectId": project["projectId"], | |
| # "created_dttm": project["created_dttm"].split('.')[0], | |
| # "description": project["description"], | |
| # "categoryName": project["categoryName"] | |
| # }) | |
| # for project in other_projects: | |
| # all_projects.append({ | |
| # "type": "Other Project", | |
| # "projectNm": project["projectNm"], | |
| # "projectId": project["projectId"], | |
| # "created_dttm": project["created_dttm"].split('.')[0], | |
| # "description": project["description"], | |
| # "categoryName": project["categoryName"] | |
| # }) | |
| # if not all_projects: | |
| # return "❌ No projects found." | |
| # # Build the formatted string | |
| # result = [f"✅ You have access to {len(all_projects)} project(s):\n"] | |
| # for i, project in enumerate(all_projects, 1): | |
| # result.append(f"{i}. Project Name: {project['projectNm']} ({project['type']})") | |
| # result.append(f" Project ID: {project['projectId']}") | |
| # result.append(f" Created On: {project['created_dttm']}") | |
| # result.append(f" Description: {project['description']}") | |
| # result.append(f" Category: {project['categoryName']}\n") | |
| # return "\n".join(result) | |
| # # === TOOL FUNCTIONS === | |
| # def search_documents(query: str) -> str: | |
| # """Search through ingested documents and get relevant information.""" | |
| # try: | |
| # # Generate embedding for the query | |
| # query_vector = embedding_model.embed_query(query) | |
| # # Search in Qdrant | |
| # search_result = qdrant_client.search( | |
| # collection_name=QDRANT_COLLECTION_NAME, | |
| # query_vector=query_vector, | |
| # limit=5, | |
| # ) | |
| # if not search_result: | |
| # return "No relevant information found in the knowledge base." | |
| # # Convert results to text content | |
| # context_texts = [] | |
| # sources = [] | |
| # for hit in search_result: | |
| # context_texts.append(hit.payload["text"]) | |
| # sources.append(hit.payload.get("source", "unknown")) | |
| # # Create a simple prompt for answering based on context | |
| # context = "\n\n".join(context_texts) | |
| # unique_sources = list(set(sources)) | |
| # # Use the LLM directly to answer the message based on context | |
| # prompt = f"""Based on the following context, answer the message: {query} | |
| # Context: | |
| # {context} | |
| # Please provide a comprehensive answer based on the context above. If the context doesn't contain enough information to answer the message, say so clearly.""" | |
| # response = llm.invoke(prompt) | |
| # return f"{response.content}\n\nSources: {', '.join(unique_sources)}" | |
| # except Exception as e: | |
| # return f"Error searching documents: {str(e)}" | |
| # # Global variables to store auth context (for tool functions) | |
| # _current_user_id = None | |
| # _current_org_id = None | |
| # _current_auth_token = None | |
| # def get_user_projects(userLoginId: str) -> str: | |
| # """Get list of projects for a user.""" | |
| # try: | |
| # # Use global auth context if available | |
| # if _current_auth_token and _current_user_id: | |
| # user_id = _current_user_id | |
| # org_id = _current_org_id or 1 | |
| # auth_token = _current_auth_token | |
| # else: | |
| # return "❌ Authentication token required. Please provide auth_token in your request." | |
| # # Encode auth token using the actual user ID and provided token | |
| # encoded_token = get_encoded_auth_token(user_id, auth_token) | |
| # # Fetch projects | |
| # data = fetch_user_projects(user_id, org_id, encoded_token) | |
| # # Format and return the project list | |
| # formatted = format_project_response(data) | |
| # return formatted | |
| # except ValueError: | |
| # return "❌ Invalid userLoginId format. Please provide a valid number." | |
| # except Exception as e: | |
| # return f"❌ Error fetching projects: {str(e)}" | |
| # def pandas_data_analysis(query_with_filepath: str) -> str: | |
| # """ | |
| # Tool for data analysis using PandasAI. | |
| # Input format: 'filepath|query' where filepath is S3 path or ufuid, and query is the analysis question. | |
| # """ | |
| # try: | |
| # # Parse the input to extract filepath and query | |
| # parts = query_with_filepath.split('|', 1) | |
| # if len(parts) != 2: | |
| # return "❌ Invalid input format. Please use: 'filepath|query' format." | |
| # filepath, query = parts | |
| # filepath = filepath.strip() | |
| # query = query.strip() | |
| # if not filepath or not query: | |
| # return "❌ Both filepath and query are required." | |
| # # Use the pandas_agent function | |
| # result = pandas_agent(filepath, query) | |
| # return result | |
| # except Exception as e: | |
| # return f"❌ Error in pandas data analysis: {str(e)}" | |
| # # === CREATE TOOLS === | |
| # document_search_tool = Tool( | |
| # name="document_search", | |
| # description="""Use this tool to search through ingested documents and get relevant information from the knowledge base. | |
| # Perfect for answering messages about uploaded documents, manuals, or any content that was previously stored. | |
| # Input should be a search query or message about the documents.""", | |
| # func=search_documents | |
| # ) | |
| # project_list_tool = Tool( | |
| # name="get_user_projects", | |
| # description="""Use this tool to get the list of projects for a user. | |
| # Perfect for when users ask about their projects, want to see available projects, or need project information. | |
| # Input should be the userLoginId (e.g., '25'). | |
| # Note: This tool requires authentication context to be set.""", | |
| # func=get_user_projects | |
| # ) | |
| # pandas_analysis_tool = Tool( | |
| # name="pandas_data_analysis", | |
| # description="""Use this tool for data analysis on CSV/Parquet files using PandasAI. | |
| # Perfect for when users ask questions about data analysis, statistics, insights, or want to query their datasets. | |
| # Input format: 'filepath|query' where: | |
| # - filepath: S3 file path (e.g., 'User-Uploaded-Raw-Files/Data2004csv1754926601269756') or ufuid (e.g., '123') | |
| # - query: Natural language question about the data (e.g., 'What are the top 5 values?', 'Show me summary statistics') | |
| # Examples: | |
| # - 'User-Uploaded-Raw-Files/mydata.csv|What is this file about?' | |
| # - '123|Show me the first 5 rows' | |
| # - 'Modified-Files/processed_data|What are the most common values in column X?' | |
| # """, | |
| # func=pandas_data_analysis | |
| # ) | |
| # # === AGENT SETUP === | |
| # def create_agent_with_session_memory(session_id: str): | |
| # """Create agent with session memory from Redis""" | |
| # # Get memory from Redis | |
| # memory_messages = get_session_memory(session_id) | |
| # agent_prompt = ChatPromptTemplate.from_messages([ | |
| # ("system", """You are a helpful AI assistant with access to multiple tools and conversation memory: | |
| # 1. **Document Search**: Search through uploaded documents and knowledge base | |
| # 2. **Project Management**: Get list of user projects and project information | |
| # 3. **Data Analysis**: Analyze CSV/Parquet files using PandasAI for insights, statistics, and queries | |
| # Your capabilities: | |
| # - Answer messages about documents using the document search tool | |
| # - Help users find their projects and project information | |
| # - Perform data analysis on uploaded datasets using natural language queries | |
| # - Remember previous conversations in this session | |
| # - Provide general assistance and information | |
| # - Use appropriate tools based on user queries | |
| # Guidelines: | |
| # - Use the document search tool when users ask about specific content, documentation, or information that might be in uploaded files | |
| # - Use the project tool when users ask about projects, want to see their projects, or need project-related information | |
| # - Use the pandas analysis tool when users ask about data analysis, statistics, insights, or want to query datasets | |
| # - For pandas analysis, you need both a filepath (S3 path or ufuid) and a query - ask for missing information if needed | |
| # - Reference previous conversation context when relevant | |
| # - Be clear about which tool you're using and what information you're providing | |
| # - If you're unsure which tool to use, you can ask for clarification | |
| # - Provide helpful, accurate, and well-formatted responses | |
| # Remember: Always use the most appropriate tool based on the user's message and conversation context to provide the best possible answer."""), | |
| # MessagesPlaceholder(variable_name="chat_history"), | |
| # ("user", "{input}"), | |
| # MessagesPlaceholder(variable_name="agent_scratchpad"), | |
| # ]) | |
| # # Create memory object | |
| # memory = ConversationBufferMemory( | |
| # memory_key="chat_history", | |
| # return_messages=True | |
| # ) | |
| # # Load existing messages into memory | |
| # for msg in memory_messages: | |
| # if msg["role"] == "user": | |
| # memory.chat_memory.add_user_message(msg["message"]) | |
| # else: | |
| # memory.chat_memory.add_ai_message(msg["message"]) | |
| # # Create tools list | |
| # tools = [document_search_tool, project_list_tool, pandas_analysis_tool] | |
| # # Create the agent | |
| # agent = create_openai_tools_agent(llm, tools, agent_prompt) | |
| # # Create the agent executor with memory | |
| # agent_executor = AgentExecutor( | |
| # agent=agent, | |
| # tools=tools, | |
| # verbose=True, | |
| # memory=memory | |
| # ) | |
| # return agent_executor, memory | |
| # # === API ENDPOINTS === | |
| # @app.post("/sessions", response_model=SessionResponse) | |
| # def create_new_session(userLoginId: int, orgId: int, auth_token: str): | |
| # """Create a new chat session""" | |
| # try: | |
| # session_data = create_session(userLoginId, orgId, auth_token) | |
| # return SessionResponse(**session_data) | |
| # except Exception as e: | |
| # raise HTTPException(status_code=500, detail=f"Error creating session: {str(e)}") | |
| # @app.get("/sessions") | |
| # def list_user_sessions(userLoginId: int): | |
| # """List all sessions for a user""" | |
| # try: | |
| # sessions = get_user_sessions(userLoginId) | |
| # return { | |
| # "userLoginId": userLoginId, | |
| # "total_sessions": len(sessions), | |
| # "sessions": sessions | |
| # } | |
| # except Exception as e: | |
| # raise HTTPException(status_code=500, detail=f"Error fetching sessions: {str(e)}") | |
| # @app.delete("/sessions/{session_id}") | |
| # def delete_user_session(session_id: str): | |
| # """Delete/close a session""" | |
| # try: | |
| # # Verify session exists | |
| # get_session(session_id) | |
| # # Delete session | |
| # delete_session(session_id) | |
| # return { | |
| # "message": f"Session {session_id} deleted successfully", | |
| # "session_id": session_id | |
| # } | |
| # except Exception as e: | |
| # raise HTTPException(status_code=500, detail=f"Error deleting session: {str(e)}") | |
| # @app.post("/bot") | |
| # def chat_with_bot(query: BotQuery): | |
| # """Main bot endpoint with session management""" | |
| # try: | |
| # # Set global auth context for tools | |
| # global _current_user_id, _current_org_id, _current_auth_token | |
| # _current_user_id = query.userLoginId | |
| # _current_org_id = query.orgId | |
| # _current_auth_token = query.auth_token | |
| # session_id = query.session_id | |
| # # Create new session if not provided | |
| # if not session_id: | |
| # session_data = create_session(query.userLoginId, query.orgId, query.auth_token) | |
| # session_id = session_data["session_id"] | |
| # else: | |
| # # Verify existing session | |
| # get_session(session_id) | |
| # # Add user message to session | |
| # user_message_id = add_message_to_session(session_id, "user", query.message) | |
| # # Create agent with session memory | |
| # agent_executor, memory = create_agent_with_session_memory(session_id) | |
| # # Use the agent to process the query | |
| # result = agent_executor.invoke({"input": query.message}) | |
| # # Add AI response to session | |
| # ai_message_id = add_message_to_session(session_id, "assistant", result["output"]) | |
| # # Update session memory in Redis | |
| # updated_messages = [] | |
| # for message in memory.chat_memory.messages: | |
| # if hasattr(message, 'content'): | |
| # role = "user" if message.__class__.__name__ == "HumanMessage" else "assistant" | |
| # updated_messages.append({ | |
| # "role": role, | |
| # "message": message.content, | |
| # "timestamp": datetime.now().isoformat() | |
| # }) | |
| # update_session_memory(session_id, updated_messages) | |
| # # Update session title after first user message | |
| # update_session_title(session_id) | |
| # # Clear auth context after use | |
| # _current_user_id = None | |
| # _current_org_id = None | |
| # _current_auth_token = None | |
| # return { | |
| # "session_id": session_id, | |
| # "user_message_id": user_message_id, | |
| # "ai_message_id": ai_message_id, | |
| # "message": query.message, | |
| # "answer": result["output"], | |
| # "userLoginId": query.userLoginId, | |
| # "agent_used": True | |
| # } | |
| # except Exception as e: | |
| # # Clear auth context on error | |
| # _current_user_id = None | |
| # _current_org_id = None | |
| # _current_auth_token = None | |
| # raise HTTPException(status_code=500, detail=f"Error processing chat: {str(e)}") | |
| # @app.get("/sessions/{session_id}/history", response_model=ChatHistoryResponse) | |
| # def get_session_history(session_id: str, n: int = QueryParam(50, description="Number of recent messages to return")): | |
| # """Get chat history for a session""" | |
| # try: | |
| # # Verify session exists | |
| # get_session(session_id) | |
| # # Get chat history | |
| # chat_data = redis_client.get(f"chat:{session_id}") | |
| # if not chat_data: | |
| # return ChatHistoryResponse( | |
| # session_id=session_id, | |
| # messages=[], | |
| # total_messages=0 | |
| # ) | |
| # messages = json.loads(chat_data) | |
| # # Get the last n messages (or all if less than n) | |
| # recent_messages = messages[-n:] if len(messages) > n else messages | |
| # # Convert to MessageResponse objects | |
| # message_responses = [MessageResponse(**msg) for msg in recent_messages] | |
| # return ChatHistoryResponse( | |
| # session_id=session_id, | |
| # messages=message_responses, | |
| # total_messages=len(messages) | |
| # ) | |
| # except Exception as e: | |
| # raise HTTPException(status_code=500, detail=f"Error fetching chat history: {str(e)}") | |
| # @app.post("/chat-documents") | |
| # def chat_documents_only(query: Query): | |
| # """Direct document search without agent""" | |
| # try: | |
| # result = search_documents(query.message) | |
| # return { | |
| # "message": query.message, | |
| # "answer": result, | |
| # "tool_used": "document_search" | |
| # } | |
| # except Exception as e: | |
| # return { | |
| # "message": query.message, | |
| # "answer": f"An error occurred: {str(e)}", | |
| # "tool_used": "document_search" | |
| # } | |
| # @app.post("/list-projects") | |
| # def list_projects(request: ProjectRequest): | |
| # """Direct project listing without agent""" | |
| # try: | |
| # # Use the provided auth token and userLoginId | |
| # encoded_token = get_encoded_auth_token(request.userLoginId, request.auth_token) | |
| # # Fetch projects | |
| # data = fetch_user_projects(request.userLoginId, request.orgId, encoded_token) | |
| # # Format and return the project list | |
| # formatted = format_project_response(data) | |
| # return { | |
| # "projects": formatted, | |
| # "tool_used": "project_list" | |
| # } | |
| # except Exception as e: | |
| # return { | |
| # "error": f"An error occurred: {str(e)}", | |
| # "tool_used": "project_list" | |
| # } | |
| # @app.post("/chat-with-pandas-agent") | |
| # def chat_with_pandas_agent(request: PandasAgentQuery): | |
| # """Direct pandas AI agent endpoint for data analysis""" | |
| # try: | |
| # result = pandas_agent(request.filepath, request.query) | |
| # return { | |
| # "filepath": request.filepath, | |
| # "query": request.query, | |
| # "answer": result, | |
| # "tool_used": "pandas_agent", | |
| # "timestamp": datetime.now().isoformat() | |
| # } | |
| # except Exception as e: | |
| # error_msg = f"An error occurred: {str(e)}" | |
| # return { | |
| # "filepath": request.filepath, | |
| # "query": request.query, | |
| # "answer": error_msg, | |
| # "tool_used": "pandas_agent", | |
| # "error": True, | |
| # "timestamp": datetime.now().isoformat() | |
| # } | |
| # @app.put("/sessions/{session_id}/title") | |
| # def refresh_session_title(session_id: str): | |
| # """Manually refresh/regenerate session title""" | |
| # try: | |
| # # Verify session exists | |
| # session_data = get_session(session_id) | |
| # # Generate new title | |
| # new_title = generate_session_title(session_id) | |
| # # Update session | |
| # session_data["title"] = new_title | |
| # redis_client.setex( | |
| # f"session:{session_id}", | |
| # 86400, # 24 hours | |
| # json.dumps(session_data) | |
| # ) | |
| # return { | |
| # "session_id": session_id, | |
| # "new_title": new_title, | |
| # "message": "Session title updated successfully" | |
| # } | |
| # except Exception as e: | |
| # raise HTTPException(status_code=500, detail=f"Error updating session title: {str(e)}") | |
| # @app.get("/redis-info") | |
| # def redis_info(): | |
| # """Get Redis connection information""" | |
| # try: | |
| # info = redis_client.info() | |
| # return { | |
| # "redis_connected": True, | |
| # "redis_version": info.get("redis_version"), | |
| # "used_memory": info.get("used_memory_human"), | |
| # "connected_clients": info.get("connected_clients"), | |
| # "total_keys": redis_client.dbsize() | |
| # } | |
| # except Exception as e: | |
| # return { | |
| # "redis_connected": False, | |
| # "error": str(e) | |
| # } | |
| # @app.get("/health") | |
| # def health(): | |
| # try: | |
| # redis_client.ping() | |
| # redis_status = "connected" | |
| # except: | |
| # redis_status = "disconnected" | |
| # return { | |
| # "status": "ok", | |
| # "tools": ["document_search", "project_list", "pandas_data_analysis"], | |
| # "agent": "active", | |
| # "session_management": "enabled", | |
| # "redis_status": redis_status, | |
| # "pandas_ai": "enabled", | |
| # "total_sessions": len(list(redis_client.scan_iter(match="session:*"))) | |
| # } | |
| # if __name__ == "__main__": | |
| # import uvicorn | |
| # try: | |
| # uvicorn.run(app, host="0.0.0.0", port=8000) | |
| # except KeyboardInterrupt: | |
| # print("\n🛑 Server stopped gracefully") | |
| # except Exception as e: | |
| # print(f"❌ Server error: {e}") | |
| # # |