Update app.py

app.py

@@ -1,9 +1,9 @@
+
 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, initialize_agent
-from langchain.agents.agent_types import AgentType
+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
@@ -27,12 +27,13 @@ import psycopg2
 from pandasai import SmartDataframe
 from pandasai.llm.openai import OpenAI as PandasOpenAI
 
-from fastapi import FastAPI, Request, UploadFile, File
+
+
+
+from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
-from qdrant_client.models import VectorParams, Distance, PointStruct, Filter, SearchRequest
-from langchain.document_loaders import PyPDFLoader, TextLoader, CSVLoader, Docx2txtLoader, BSHTMLLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-import tempfile
+import json
+
 
 # Import your existing S3 connection details
 from retrive_secrects import *  # CONNECTIONS_HOST, etc.
@@ -42,20 +43,10 @@ warnings.filterwarnings("ignore", message="Qdrant client version.*is incompatibl
 
 load_dotenv()
 
-app = FastAPI(title="AI Agent with Enhanced Tools and Session Management")
-
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-if not OPENAI_API_KEY:
-    raise ValueError("❌ OPENAI_API_KEY not set in environment variables")
-
-# Qdrant Configuration
-API_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhY2Nlc3MiOiJtIiwiZXhwIjoxNzY0MTQ5OTc3fQ.l_2R-Eyb_530887EGLUkawZQamhPGVklDMlaVs0bDqo"
-URL = "https://09476415-f871-4664-9c92-2f7f17c223ee.eu-central-1-0.aws.cloud.qdrant.io"
-
-# Initialize Qdrant client
-client = QdrantClient(url=URL, api_key=API_KEY)
+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))
@@ -66,10 +57,24 @@ 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
+# 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
@@ -77,16 +82,19 @@ 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,
@@ -95,6 +103,7 @@ def get_redis_client():
                 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
@@ -102,23 +111,35 @@ def get_redis_client():
         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",#3072
+    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)
 
-# === ENHANCED DATA READING FUNCTIONS ===
+# === PANDAS AI FUNCTIONS ===
 
 def read_parquet_file_from_s3(ufuid=None, columns_list=None, records_count=None, file_location=''):
     """
-    Enhanced version that reads Parquet files from S3 using Dask and returns as Pandas DataFrame.
+    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,
@@ -135,18 +156,22 @@ def read_parquet_file_from_s3(ufuid=None, columns_list=None, records_count=None,
                 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",
@@ -154,9 +179,10 @@ def read_parquet_file_from_s3(ufuid=None, columns_list=None, records_count=None,
             assume_missing=True
         )
 
-        ddf = ddf.repartition(npartitions=8)
-        print("✅ Reading Parquet file from S3 completed successfully.")
+        ddf = ddf.repartition(npartitions=8)  # Optimize for processing
+        print("Reading Parquet file from S3 completed successfully.")
         
+        # Close database connection
         cursor.close()
         conn.close()
         
@@ -164,142 +190,56 @@ def read_parquet_file_from_s3(ufuid=None, columns_list=None, records_count=None,
 
     except Exception as e:
         print(f"❌ Error reading Parquet file: {e}")
-        return pd.DataFrame()
-
-# === ENHANCED TOOL CLASSES ===
+        return pd.DataFrame()  # Return empty DataFrame on error
 
-class PandasAITool:
-    """Enhanced PandasAI tool with better error handling and context"""
+def pandas_agent(filepath: str, query: str) -> str:
+    """
+    PandasAI agent that reads data from S3 and answers queries about the data.
     
-    def __init__(self, df: pd.DataFrame = None):
-        if not OPENAI_API_KEY:
-            raise ValueError("OPENAI_API_KEY is not set in environment variables.")
-        
-        self.llm = PandasOpenAI(api_token=OPENAI_API_KEY)
-        self.df = df
-        self.sdf = SmartDataframe(df, config={"llm": self.llm}) if df is not None else None
-
-    def run(self, query_with_filepath: str) -> str:
-        """
-        Enhanced run method that can handle both direct queries and file loading
-        Input format: 'filepath|query' or just 'query' if DataFrame already loaded
-        """
-        try:
-            # Parse input to extract filepath and query if both provided
-            if '|' in query_with_filepath:
-                parts = query_with_filepath.split('|', 1)
-                if len(parts) == 2:
-                    filepath, query = parts
-                    filepath = filepath.strip()
-                    query = query.strip()
-                    
-                    # Load data from file
-                    if filepath.isdigit():
-                        data = read_parquet_file_from_s3(ufuid=int(filepath))
-                    else:
-                        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."
-                    
-                    # Update the dataframe and SmartDataframe
-                    self.df = data
-                    self.sdf = SmartDataframe(data, config={"llm": self.llm})
-                else:
-                    query = query_with_filepath
-            else:
-                query = query_with_filepath
-
-            if self.sdf is None:
-                return "❌ No dataset loaded. Please provide a file path with your query."
-
-            print(f"\n[PandasAI Query]: {query}")
-            result = self.sdf.chat(query)
-            print(f"[PandasAI Result]: {result}")
-            
-            # 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 PandasAI Tool: {str(e)}"
-            print(error_msg)
-            return error_msg
+    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."
 
-class EnhancedDocumentSearchTool:
-    """Enhanced document search tool that works with both local and remote Qdrant"""
-    
-    def __init__(self, collection_name: str = None):
-        self.collection_name = collection_name or QDRANT_COLLECTION_NAME
-        self.client = client
-        self.embeddings = embedding_model
-        self.llm = llm
-
-    def run(self, query: str) -> str:
-        """
-        Enhanced document search with better context handling
-        Can handle: 'search query' or 'ufuid search query'
-        """
-        try:
-            # Check if query starts with a ufuid
-            parts = query.strip().split(maxsplit=1)
-            if len(parts) >= 2 and parts[0].isdigit():
-                ufuid = int(parts[0])
-                search_query = parts[1]
-                # Could potentially filter by ufuid if metadata supports it
-            else:
-                search_query = query
-
-            print(f"[DocumentSearchTool] Searching for: {search_query}")
-
-            # Generate embedding for the query
-            query_vector = self.embeddings.embed_query(search_query)
-            
-            # Search in Qdrant
-            search_result = self.client.search(
-                collection_name=self.collection_name,
-                query_vector=query_vector,
-                limit=5,
-            )
-            
-            if not search_result:
-                return "❌ No relevant information found in the knowledge base."
-            
-            # Extract context and sources
-            context_texts = []
-            sources = []
-            
-            for hit in search_result:
-                context_texts.append(hit.payload.get("text", str(hit.payload)))
-                sources.append(hit.payload.get("source", "unknown"))
-            
-            context = "\n\n".join(context_texts)
-            unique_sources = list(set(sources))
-            
-            # Use LLM to provide a comprehensive answer
-            prompt = f"""Based on the following context from documents, answer this query: {search_query}
+        # Initialize PandasAI LLM
+        if not OPENAI_API_KEY:
+            return "❌ OPENAI_API_KEY is not set in environment variables."
 
-Context from documents:
-{context}
+        pandas_llm = PandasOpenAI(api_token=OPENAI_API_KEY)
 
-Please provide a comprehensive and helpful answer based on the context above. If the context doesn't contain enough information to fully answer the query, mention this clearly but provide what information is available."""
+        # Create SmartDataframe
+        sdf = SmartDataframe(data, config={"llm": pandas_llm})
 
-            response = self.llm.invoke(prompt)
-            
-            result = f"{response.content}\n\n📚 Sources: {', '.join(unique_sources)}"
-            print(f"[DocumentSearchTool Result]: {result}")
-            return result
-            
-        except Exception as e:
-            error_msg = f"❌ Error in Document Search Tool: {str(e)}"
-            print(error_msg)
-            return error_msg
+        # 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 ===
 
@@ -318,6 +258,10 @@ class BotQuery(BaseModel):
     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
@@ -329,7 +273,7 @@ class SessionResponse(BaseModel):
 class MessageResponse(BaseModel):
     message_id: str
     session_id: str
-    role: str
+    role: str  # "user" or "assistant"
     message: str
     timestamp: str
 
@@ -344,112 +288,34 @@ class UpdateSessionTitleRequest(BaseModel):
 # === SESSION MANAGEMENT FUNCTIONS ===
 
 
-def get_user_sessions(userLoginId: int) -> List[dict]:
-    """Get all sessions for a user - returns stored titles, not generated ones"""
-    sessions = []
-    
-    # Scan for all session keys
-    for key in redis_client.scan_iter(match="session:*"):
-        session_data = redis_client.get(key)
-        if session_data:
-            try:
-                session = json.loads(session_data)
-                if session["userLoginId"] == userLoginId:
-                    # Use the stored title, don't override it
-                    # Only generate title if it's missing or still "New Chat"
-                    current_title = session.get("title", "New Chat")
-                    
-                    # Only auto-generate if title is missing or default
-                    if not current_title or current_title == "New Chat":
-                        session["title"] = generate_session_title(session["session_id"])
-                        # Update the session in Redis with the new title
-                        redis_client.setex(
-                            key,
-                            86400,  # 24 hours
-                            json.dumps(session)
-                        )
-                    
-                    sessions.append(session)
-            except json.JSONDecodeError:
-                # Skip corrupted session data
-                continue
-    
-    # Sort sessions by created_at (most recent first)
-    sessions.sort(key=lambda x: x.get("created_at", ""), reverse=True)
-    return sessions
 
 
-def update_session_title_endpoint(session_id: str, request: UpdateSessionTitleRequest):
-    """Update the title of an existing session with custom name"""
-    try:
-        # Get current session data from Redis
-        redis_key = f"session:{session_id}"
-        session_data = redis_client.get(redis_key)
-        
-        if not session_data:
-            raise HTTPException(status_code=404, detail="Session not found or expired")
-        
-        # Parse current session data
-        session = json.loads(session_data)
-        
-        # Validate new title
-        new_title = request.new_title.strip()
-        if not new_title:
-            raise HTTPException(status_code=400, detail="New title cannot be empty")
-        
-        if len(new_title) > 100:
-            raise HTTPException(status_code=400, detail="Title cannot exceed 100 characters")
-        
-        # Get the actual old title from Redis (not from generate_session_title)
-        old_title = session.get("title", "New Chat")
-        
-        # Update the title
-        session["title"] = new_title
-        session["last_updated"] = datetime.now().isoformat()
-        
-        # Save updated session back to Redis with same TTL
-        redis_client.setex(
-            redis_key,
-            86400,  # 24 hours TTL
-            json.dumps(session)
-        )
-        
-        # Verify the update by reading back from Redis
-        updated_session_data = redis_client.get(redis_key)
-        if updated_session_data:
-            updated_session = json.loads(updated_session_data)
-            actual_new_title = updated_session.get("title", "Unknown")
-        else:
-            raise HTTPException(status_code=500, detail="Failed to verify session update")
-        
-        print(f"Title update: {old_title} -> {actual_new_title}")  # Debug log
-        
-        return {
-            "message": "Session title updated successfully",
-            "session_id": session_id,
-            "old_title": old_title,
-            "new_title": actual_new_title,
-            "last_updated": session.get("last_updated"),
-            "success": True
-        }
-        
-    except HTTPException:
-        raise
-    except json.JSONDecodeError:
-        raise HTTPException(status_code=500, detail="Invalid session data format in Redis")
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error updating session title: {str(e)}")
+@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
+
 
 
-# Also update the UpdateSessionTitleRequest model to be more explicit
-class UpdateSessionTitleRequest(BaseModel):
-    new_title: str = Field(
-        ..., 
-        min_length=1, 
-        max_length=100, 
-        description="New title for the session",
-        example="My Custom Chat Title"
-    )
 
 def create_session(userLoginId: int, orgId: int, auth_token: str) -> dict:
     """Create a new chat session"""
@@ -461,12 +327,29 @@ def create_session(userLoginId: int, orgId: int, auth_token: str) -> dict:
         "auth_token": auth_token,
         "created_at": datetime.now().isoformat(),
         "status": "active",
-        "title": "New Chat"
+        "title": "New Chat"  # Default title, will be updated after first message
     }
     
-    redis_client.setex(f"session:{session_id}", 86400, json.dumps(session_data))
-    redis_client.setex(f"chat:{session_id}", 86400, json.dumps([]))
-    redis_client.setex(f"memory:{session_id}", 86400, json.dumps([]))
+    # 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
 
@@ -488,25 +371,70 @@ def add_message_to_session(session_id: str, role: str, message: str) -> str:
         "timestamp": datetime.now().isoformat()
     }
     
+    # Get current chat history
     chat_history = redis_client.get(f"chat:{session_id}")
-    messages = json.loads(chat_history) if chat_history else []
+    if chat_history:
+        messages = json.loads(chat_history)
+    else:
+        messages = []
+    
+    # Add new message
     messages.append(message_data)
     
-    redis_client.setex(f"chat:{session_id}", 86400, json.dumps(messages))
+    # 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}")
-    return json.loads(memory_data) if memory_data else []
+    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, json.dumps(messages))
+    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"
@@ -515,6 +443,7 @@ def generate_session_title(session_id: str) -> str:
         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":
@@ -524,41 +453,59 @@ def generate_session_title(session_id: str) -> str:
         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."""
+Return only the title, no quotes or additional text. The title should capture the main topic or intent of the conversation."""
 
-        response = llm.invoke(title_prompt)
-        title = response.content.strip().replace('"', '').replace("'", "")
-        
-        if len(title) > 50:
-            title = title[:47] + "..."
-        
-        return title if title else "New Chat"
+        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 update_session_title(session_id: str):
-    """Update session title after first message"""
-    try:
-        session_data = redis_client.get(f"session:{session_id}")
-        if not session_data:
-            return
-        
-        session = json.loads(session_data)
-        
-        if session.get("title", "New Chat") == "New Chat":
-            new_title = generate_session_title(session_id)
-            session["title"] = new_title
-            
-            redis_client.setex(f"session:{session_id}", 86400, json.dumps(session))
-            
-    except Exception as e:
-        print(f"Error updating session title: {e}")
+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 ===
 
@@ -568,7 +515,11 @@ def get_encoded_auth_token(user: int, token: str) -> str:
 
 def fetch_user_projects(userLoginId: int, orgId: int, auth_token: str):
     url = "https://japidemo.dev.ingenspark.com/fetchUserProjects"
-    payload = {"userLoginId": userLoginId, "orgId": orgId}
+    payload = {
+        "userLoginId": userLoginId,
+        "orgId": orgId
+    }
+    
     headers = {
         'accept': 'application/json, text/plain, */*',
         'authorization': f'Basic {auth_token}',
@@ -612,6 +563,7 @@ def format_project_response(data: dict) -> str:
     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']})")
@@ -621,162 +573,254 @@ def format_project_response(data: dict) -> str:
         result.append(f"   Category: {project['categoryName']}\n")
     return "\n".join(result)
 
-# Global variables for auth context
+# === 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"""
+    """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."
+            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)
-        return format_project_response(data)
         
+        # 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 ENHANCED AGENT ===
+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
+)
 
-def create_enhanced_agent_with_session_memory(session_id: str):
-    """Create enhanced agent with all tools and session memory"""
+# === 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)
     
-    # Initialize tools with enhanced capabilities
-    pandas_tool = PandasAITool()
-    doc_search_tool = EnhancedDocumentSearchTool()
-    
-    # Create tool wrappers
-    tools = [
-        Tool(
-            name="enhanced_pandas_analysis",
-            description="""Use this tool for advanced data analysis on CSV/Parquet files using PandasAI.
-            Input format: 'filepath|query' where:
-            - filepath: S3 file path or ufuid (e.g., 'User-Uploaded-Raw-Files/data.csv' or '123')
-            - query: Natural language question about the data
-            
-            Examples:
-            - 'User-Uploaded-Raw-Files/mydata.csv|What are the top 5 values?'
-            - '123|Show me summary statistics'
-            - 'Modified-Files/processed_data|What are the trends in sales data?'""",
-            func=pandas_tool.run
-        ),
-        Tool(
-            name="enhanced_document_search",
-            description="""Use this tool to search through ingested documents and knowledge base.
-            Input can be:
-            - Simple search query: 'search terms'
-            - With ufuid: 'ufuid search terms' (e.g., '9 list all files')
-            
-            Perfect for answering questions about uploaded documents, manuals, or stored content.""",
-            func=doc_search_tool.run
-        ),
-        Tool(
-            name="get_user_projects",
-            description="""Use this tool to get the list of projects for a user.
-            Input should be the userLoginId as a string (e.g., '25').
-            Perfect for when users ask about their projects or need project information.""",
-            func=get_user_projects
-        )
-    ]
-    
-    # Create agent prompt
     agent_prompt = ChatPromptTemplate.from_messages([
-        ("system", """You are an advanced AI assistant with access to powerful data analysis and document search capabilities:
-
-🔧 **Available Tools:**
-1. **Enhanced Pandas Analysis**: Analyze CSV/Parquet files with natural language queries
-2. **Enhanced Document Search**: Search through uploaded documents and knowledge base  
-3. **Project Management**: Get user project information
-
-💡 **Your Capabilities:**
-- Perform sophisticated data analysis on datasets using PandasAI
-- Search and retrieve information from document knowledge bases
-- Help users manage and find their projects
-- Remember conversation context and provide personalized assistance
-- Handle complex multi-step queries involving multiple tools
-
-📋 **Guidelines:**
-- Use pandas analysis when users ask about data insights, statistics, trends, or file analysis
-- Use document search when users ask about specific content, documentation, or information
-- Use project tools when users need project-related information
-- Always provide clear, comprehensive, and well-formatted responses
+        ("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
-- If you need clarification, ask specific questions to help the user better
+- 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
 
-🎯 **Response Quality:**
-- Be thorough but concise
-- Use proper formatting for readability
-- Explain your reasoning when using tools
-- Provide actionable insights and next steps when appropriate
-
-Remember: You have access to both the user's conversation history and powerful analysis tools to provide the best possible assistance."""),
+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 and load existing messages
+    # 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 the enhanced agent
+    # 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 agent executor with memory
+    # Create the agent executor with memory
     agent_executor = AgentExecutor(
         agent=agent, 
         tools=tools, 
         verbose=True,
-        memory=memory,
-        max_iterations=5,
-        early_stopping_method="generate"
+        memory=memory
     )
     
     return agent_executor, memory
 
-# === MIDDLEWARE ===
+# === API ENDPOINTS ===
 
-@app.middleware("http")
-async def add_success_flag(request: Request, call_next):
-    response = await call_next(request)
+@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)}")
 
-    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())
-            data["success"] = 200 <= response.status_code < 300
-            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:
-            pass
-    return response
+@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)}")
 
-# === MAIN BOT ENDPOINT ===
+@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 enhanced_chat_with_bot(query: BotQuery):
-    """Enhanced main bot endpoint with combined tools and session management"""
+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
@@ -797,13 +841,11 @@ def enhanced_chat_with_bot(query: BotQuery):
         # Add user message to session
         user_message_id = add_message_to_session(session_id, "user", query.message)
         
-        # Create enhanced agent with session memory
-        agent_executor, memory = create_enhanced_agent_with_session_memory(session_id)
+        # Create agent with session memory
+        agent_executor, memory = create_agent_with_session_memory(session_id)
         
-        # Process the query with enhanced agent
-        print(f"\n🤖 Processing query: {query.message}")
+        # Use the agent to process the query
         result = agent_executor.invoke({"input": query.message})
-        print(f"🎯 Agent response: {result['output']}")
         
         # Add AI response to session
         ai_message_id = add_message_to_session(session_id, "assistant", result["output"])
@@ -836,9 +878,7 @@ def enhanced_chat_with_bot(query: BotQuery):
             "message": query.message,
             "answer": result["output"],
             "userLoginId": query.userLoginId,
-            "agent_type": "enhanced_multi_tool",
-            "tools_available": ["enhanced_pandas_analysis", "enhanced_document_search", "get_user_projects"],
-            "timestamp": datetime.now().isoformat()
+            "agent_used": True
         }
         
     except Exception as e:
@@ -847,65 +887,16 @@ def enhanced_chat_with_bot(query: BotQuery):
         _current_org_id = None
         _current_auth_token = None
         
-        error_msg = f"❌ Error processing enhanced chat: {str(e)}"
-        print(error_msg)
-        raise HTTPException(status_code=500, detail=error_msg)
-
-# === SESSION MANAGEMENT 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 = []
-        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:
-                    session["title"] = generate_session_title(session["session_id"])
-                    sessions.append(session)
-        
-        sessions.sort(key=lambda x: x["created_at"], reverse=True)
-        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:
-        get_session(session_id)
-        redis_client.delete(f"session:{session_id}")
-        redis_client.delete(f"chat:{session_id}")
-        redis_client.delete(f"memory:{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)}")
+        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")):
+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(
@@ -915,7 +906,11 @@ def get_session_history(session_id: str, n: int = QueryParam(50, description="Nu
             )
         
         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(
@@ -926,49 +921,19 @@ def get_session_history(session_id: str, n: int = QueryParam(50, description="Nu
         
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error fetching chat history: {str(e)}")
+    
+
 
-# @app.put("/sessions/{session_id}/title")
-# def update_session_title_endpoint(session_id: str, request: UpdateSessionTitleRequest):
-#     """Update the title of an existing session"""
-#     try:
-#         session_data = redis_client.get(f"session:{session_id}")
-#         if not session_data:
-#             raise HTTPException(status_code=404, detail="Session not found or expired")
-        
-#         session = json.loads(session_data)
-        
-#         new_title = request.new_title.strip()
-#         if not new_title:
-#             raise HTTPException(status_code=400, detail="New title cannot be empty")
-        
-#         if len(new_title) > 100:
-#             raise HTTPException(status_code=400, detail="Title cannot exceed 100 characters")
-        
-#         old_title = session.get("title", "New Chat")
-#         session["title"] = new_title
-#         session["last_updated"] = datetime.now().isoformat()
-        
-#         redis_client.setex(f"session:{session_id}", 86400, json.dumps(session))
-        
-#         return {
-#             "message": "Session title updated successfully",
-#             "session_id": session_id,
-#             "old_title": old_title,
-#             "new_title": new_title
-#         }
-        
-#     except HTTPException:
-#         raise
-#     except Exception as e:
-#         raise HTTPException(status_code=500, detail=f"Error updating session title: {str(e)}")
 @app.put("/sessions/{session_id}/title")
 def update_session_title_endpoint(session_id: str, request: UpdateSessionTitleRequest):
-    """Update the title of an existing session with custom name"""
+    """Update the title of an existing session"""
     try:
-        # Get current session data from Redis
-        redis_key = f"session:{session_id}"
-        session_data = redis_client.get(redis_key)
+        # Verify that the session_id in URL matches the one in request body
+        if session_id != request.session_id:
+            raise HTTPException(status_code=400, detail="Session ID in URL and request body must match")
         
+        # Verify session exists and get current session data
+        session_data = redis_client.get(f"session:{session_id}")
         if not session_data:
             raise HTTPException(status_code=404, detail="Session not found or expired")
         
@@ -984,302 +949,215 @@ def update_session_title_endpoint(session_id: str, request: UpdateSessionTitleRe
             raise HTTPException(status_code=400, detail="Title cannot exceed 100 characters")
         
         # Update the title
-        old_title = session.get("title", "New Chat")
+        old_title = session.get("title", "Unknown")
         session["title"] = new_title
         session["last_updated"] = datetime.now().isoformat()
         
-        # Save updated session back to Redis with same TTL
+        # Save updated session back to Redis
         redis_client.setex(
-            redis_key,
+            f"session:{session_id}",
             86400,  # 24 hours TTL
             json.dumps(session)
         )
         
-        # Verify the update by reading back from Redis
-        updated_session_data = redis_client.get(redis_key)
-        if updated_session_data:
-            updated_session = json.loads(updated_session_data)
-            actual_new_title = updated_session.get("title", "Unknown")
-        else:
-            raise HTTPException(status_code=500, detail="Failed to verify session update")
-        
         return {
             "message": "Session title updated successfully",
             "session_id": session_id,
             "old_title": old_title,
-            "new_title": actual_new_title,
-            "last_updated": session.get("last_updated"),
-            "success": True
+            "new_title": new_title
         }
         
     except HTTPException:
         raise
-    except json.JSONDecodeError:
-        raise HTTPException(status_code=500, detail="Invalid session data format in Redis")
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error updating session title: {str(e)}")
 
-
-# Alternative endpoint for auto-generating titles based on chat history
-@app.put("/sessions/{session_id}/title/auto-generate")
-def auto_generate_session_title(session_id: str):
-    """Automatically generate and update session title based on chat history"""
-    try:
-        # Verify session exists
-        redis_key = f"session:{session_id}"
-        session_data = redis_client.get(redis_key)
-        
-        if not session_data:
-            raise HTTPException(status_code=404, detail="Session not found or expired")
-        
-        session = json.loads(session_data)
-        old_title = session.get("title", "New Chat")
-        
-        # Generate new title based on chat history
-        new_title = generate_session_title(session_id)
-        
-        # Update session
-        session["title"] = new_title
-        session["last_updated"] = datetime.now().isoformat()
-        
-        # Save updated session back to Redis
-        redis_client.setex(
-            redis_key,
-            86400,  # 24 hours TTL
-            json.dumps(session)
-        )
-        
-        return {
-            "message": "Session title auto-generated successfully",
-            "session_id": session_id,
-            "old_title": old_title,
-            "new_title": new_title,
-            "method": "auto_generated",
-            "success": True
-        }
-        
-    except HTTPException:
-        raise
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error auto-generating session title: {str(e)}")
-# === DIRECT TOOL ENDPOINTS ===
-
 @app.post("/chat-documents")
 def chat_documents_only(query: Query):
-    """Direct enhanced document search without agent"""
+    """Direct document search without agent"""
     try:
-        doc_search_tool = EnhancedDocumentSearchTool()
-        result = doc_search_tool.run(query.message)
+        result = search_documents(query.message)
         return {
             "message": query.message,
             "answer": result,
-            "tool_used": "enhanced_document_search",
-            "timestamp": datetime.now().isoformat()
+            "tool_used": "document_search"
         }
     except Exception as e:
         return {
             "message": query.message,
-            "answer": f"Error occurred: {str(e)}",
-            "tool_used": "enhanced_document_search",
-            "error": True,
-            "timestamp": datetime.now().isoformat()
+            "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",
-            "timestamp": datetime.now().isoformat()
+            "tool_used": "project_list"
         }
     except Exception as e:
         return {
-            "error": f"Error occurred: {str(e)}",
-            "tool_used": "project_list",
-            "timestamp": datetime.now().isoformat()
+            "error": f"An error occurred: {str(e)}",
+            "tool_used": "project_list"
         }
 
 @app.post("/chat-with-pandas-agent")
-def chat_with_enhanced_pandas_agent(request: Query):
-    """Direct enhanced pandas AI agent endpoint for data analysis"""
+def chat_with_pandas_agent(request: PandasAgentQuery):
+    """Direct pandas AI agent endpoint for data analysis"""
     try:
-        pandas_tool = PandasAITool()
-        result = pandas_tool.run(request.message)
+        result = pandas_agent(request.filepath, request.query)
         
         return {
-            "query": request.message,
+            "filepath": request.filepath,
+            "query": request.query,
             "answer": result,
-            "tool_used": "enhanced_pandas_analysis",
+            "tool_used": "pandas_agent",
             "timestamp": datetime.now().isoformat()
         }
         
     except Exception as e:
-        error_msg = f"Error occurred: {str(e)}"
+        error_msg = f"An error occurred: {str(e)}"
         return {
-            "query": request.message,
+            "filepath": request.filepath,
+            "query": request.query,
             "answer": error_msg,
-            "tool_used": "enhanced_pandas_analysis",
+            "tool_used": "pandas_agent",
             "error": True,
             "timestamp": datetime.now().isoformat()
         }
 
-# === COLLECTION MANAGEMENT ENDPOINTS ===
-
-class CollectionRequest(BaseModel):
-    name: str
-    vector_size: int
-    distance: str = "Cosine"
-
-class UpdateCollectionRequest(BaseModel):
-    vector_size: int = None
-    distance: str = None
-
-@app.post("/collections/")
-def create_collection(req: CollectionRequest):
-    """Create a new Qdrant collection"""
-    distance_map = {
-        "Cosine": Distance.COSINE,
-        "Euclid": Distance.EUCLID,
-        "Dot": Distance.DOT,
-    }
-    if req.distance not in distance_map:
-        raise HTTPException(status_code=400, detail="Invalid distance metric")
-
+@app.put("/sessions/{session_id}/title")
+def refresh_session_title(session_id: str):
+    """Manually refresh/regenerate session title"""
     try:
-        client.recreate_collection(
-            collection_name=req.name,
-            vectors_config=VectorParams(size=req.vector_size, distance=distance_map[req.distance]),
+        # 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 {"message": f"Collection '{req.name}' created successfully"}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.get("/collections/")
-def list_collections():
-    """List all Qdrant collections"""
-    try:
-        collections = client.get_collections()
-        return collections.dict()
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.get("/collections/{name}")
-def get_collection(name: str):
-    """Get collection information"""
-    try:
-        collection_info = client.get_collection(collection_name=name)
-        return collection_info.dict()
+        
+        return {
+            "session_id": session_id,
+            "new_title": new_title,
+            "message": "Session title updated successfully"
+        }
+        
     except Exception as e:
-        raise HTTPException(status_code=404, detail=f"Collection '{name}' not found: {str(e)}")
-
-@app.delete("/collections/{name}")
-def delete_collection(name: str):
-    """Delete a collection"""
+        raise HTTPException(status_code=500, detail=f"Error updating session title: {str(e)}")
+    
+@app.put("/sessions/{session_id}/title")
+def update_session_title_endpoint(session_id: str, request: UpdateSessionTitleRequest):
+    """Update the title of an existing session with custom name"""
     try:
-        client.delete_collection(collection_name=name)
-        return {"message": f"Collection '{name}' deleted successfully"}
+        # Verify session exists and get current session data
+        session_data = redis_client.get(f"session:{session_id}")
+        if not session_data:
+            raise HTTPException(status_code=404, detail="Session not found or expired")
+        
+        # Parse current session data
+        session = json.loads(session_data)
+        
+        # Validate new title
+        new_title = request.new_title.strip()
+        if not new_title:
+            raise HTTPException(status_code=400, detail="New title cannot be empty")
+        
+        if len(new_title) > 100:
+            raise HTTPException(status_code=400, detail="Title cannot exceed 100 characters")
+        
+        # Update the title
+        old_title = session.get("title", "New Chat")
+        session["title"] = new_title
+        session["last_updated"] = datetime.now().isoformat()
+        
+        # Save updated session back to Redis
+        redis_client.setex(
+            f"session:{session_id}",
+            86400,  # 24 hours TTL
+            json.dumps(session)
+        )
+        
+        return {
+            "message": "Session title updated successfully",
+            "session_id": session_id,
+            "old_title": old_title,
+            "new_title": new_title
+        }
+        
+    except HTTPException:
+        raise
     except Exception as e:
-        raise HTTPException(status_code=404, detail=f"Collection '{name}' not found: {str(e)}")
-
-# === DOCUMENT INGESTION ENDPOINTS ===
-
-@app.post("/ingest/{collection_name}")
-async def ingest_file(collection_name: str, file: UploadFile = File(...)):
-    """Ingest documents into Qdrant collection"""
-    suffix = os.path.splitext(file.filename)[-1].lower()
-    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
-        tmp.write(await file.read())
-        tmp_path = tmp.name
-
-    try:
-        # Select loader based on file suffix
-        if suffix == ".pdf":
-            loader = PyPDFLoader(tmp_path)
-        elif suffix in [".txt", ".md"]:
-            loader = TextLoader(tmp_path)
-        elif suffix == ".csv":
-            loader = CSVLoader(file_path=tmp_path)
-        elif suffix == ".docx":
-            loader = Docx2txtLoader(tmp_path)
-        elif suffix == ".html":
-            loader = BSHTMLLoader(file_path=tmp_path)
-        else:
-            raise HTTPException(status_code=400, detail=f"Unsupported file type: {suffix}")
-
-        docs = loader.load()
-        splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-        chunks = splitter.split_documents(docs)
-
-        texts = [chunk.page_content for chunk in chunks]
-        embeddings = embedding_model.embed_documents(texts)
-
-        # Verify embedding dimension
-        collection_info = client.get_collection(collection_name=collection_name)
-        expected_dim = collection_info.config.params.vectors.size
-        if len(embeddings[0]) != expected_dim:
-            raise HTTPException(
-                status_code=400,
-                detail=f"Embedding dimension mismatch: expected {expected_dim}, got {len(embeddings[0])}",
-            )
-
-        points = [
-            PointStruct(
-                id=str(uuid.uuid4()),
-                vector=embeddings[i],
-                payload={"text": texts[i], "source": file.filename},
-            )
-            for i in range(len(texts))
-        ]
-
-        client.upsert(collection_name=collection_name, points=points)
+        raise HTTPException(status_code=500, detail=f"Error updating session title: {str(e)}")# Add this to your existing Pydantic models section
+class UpdateSessionTitleRequest(BaseModel):
+    session_id: str
+    new_title: str
 
-    except HTTPException as he:
-        raise he
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Ingestion failed: {str(e)}")
-    finally:
-        os.remove(tmp_path)
-
-    return {"message": f"'{file.filename}' ingested into '{collection_name}' successfully"}
-
-@app.get("/search/{collection_name}")
-def search_collection(
-    collection_name: str,
-    query: str = QueryParam(..., description="Search query"),
-    top_k: int = 5
-):
-    """Search documents in collection"""
+# Add this endpoint to your FastAPI app
+@app.put("/sessions/{session_id}/title")
+def update_session_title_endpoint(session_id: str, request: UpdateSessionTitleRequest):
+    """Update the title of an existing session"""
     try:
-        query_vector = embedding_model.embed_query(query)
-        search_result = client.search(
-            collection_name=collection_name,
-            query_vector=query_vector,
-            limit=top_k,
+        # Verify that the session_id in URL matches the one in request body
+        if session_id != request.session_id:
+            raise HTTPException(status_code=400, detail="Session ID in URL and request body must match")
+        
+        # Verify session exists and get current session data
+        session_data = redis_client.get(f"session:{session_id}")
+        if not session_data:
+            raise HTTPException(status_code=404, detail="Session not found or expired")
+        
+        # Parse current session data
+        session = json.loads(session_data)
+        
+        # Validate new title
+        new_title = request.new_title.strip()
+        if not new_title:
+            raise HTTPException(status_code=400, detail="New title cannot be empty")
+        
+        if len(new_title) > 100:
+            raise HTTPException(status_code=400, detail="Title cannot exceed 100 characters")
+        
+        # Update the title
+        old_title = session.get("title", "Unknown")
+        session["title"] = new_title
+        session["last_updated"] = datetime.now().isoformat()
+        
+        # Save updated session back to Redis
+        redis_client.setex(
+            f"session:{session_id}",
+            86400,  # 24 hours TTL
+            json.dumps(session)
         )
-
-        results = [
-            {
-                "score": hit.score,
-                "payload": hit.payload,
-            }
-            for hit in search_result
-        ]
-
+        
         return {
-            "query": query,
-            "results": results,
-            "timestamp": datetime.now().isoformat()
+            "message": "Session title updated successfully",
+            "session_id": session_id,
+            "old_title": old_title,
+            "new_title": new_title
         }
-
+        
+    except HTTPException:
+        raise
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}")
-
-# === HEALTH AND INFO ENDPOINTS ===
+        raise HTTPException(status_code=500, detail=f"Error updating session title: {str(e)}")
 
 @app.get("/redis-info")
 def redis_info():
@@ -1301,63 +1179,29 @@ def redis_info():
 
 @app.get("/health")
 def health():
-    """Health check endpoint"""
     try:
         redis_client.ping()
         redis_status = "connected"
     except:
         redis_status = "disconnected"
     
-    try:
-        client.get_collections()
-        qdrant_status = "connected"
-    except:
-        qdrant_status = "disconnected"
-    
     return {
         "status": "ok",
-        "tools": ["enhanced_pandas_analysis", "enhanced_document_search", "get_user_projects"], 
-        "agent": "enhanced_multi_tool",
+        "tools": ["document_search", "project_list", "pandas_data_analysis"], 
+        "agent": "active",
         "session_management": "enabled",
         "redis_status": redis_status,
-        "qdrant_status": qdrant_status,
-        "pandas_ai": "enhanced",
-        "total_sessions": len(list(redis_client.scan_iter(match="session:*"))),
-        "timestamp": datetime.now().isoformat()
-    }
-
-@app.get("/")
-def root():
-    """Root endpoint with API information"""
-    return {
-        "message": "Enhanced AI Agent API with Multi-Tool Support",
-        "version": "2.0",
-        "features": [
-            "Enhanced PandasAI data analysis",
-            "Advanced document search with Qdrant",
-            "Project management integration",
-            "Session-based conversation memory",
-            "Redis-backed session storage",
-            "Multi-tool agent with context awareness"
-        ],
-        "endpoints": {
-            "chat": "/bot",
-            "sessions": "/sessions",
-            "direct_tools": ["/chat-documents", "/list-projects", "/chat-with-pandas-agent"],
-            "collections": "/collections/",
-            "health": "/health"
-        },
-        "timestamp": datetime.now().isoformat()
+        "pandas_ai": "enabled",
+        "total_sessions": len(list(redis_client.scan_iter(match="session:*")))
     }
 
 if __name__ == "__main__":
     import uvicorn
     try:
-        print("Starting Enhanced AI Agent API...")
-        print("Features: Enhanced PandasAI, Document Search, Project Management")
-        print("Session Management: Redis-backed with conversation memory")
         uvicorn.run(app, host="0.0.0.0", port=8000)
     except KeyboardInterrupt:
-        print("\nServer stopped gracefully")
+        print("\n🛑 Server stopped gracefully")
     except Exception as e:
-        print(f"Server error: {e}")
+        print(f"❌ Server error: {e}")
+
+#bot4