Upload 8 files

.gitignore

@@ -0,0 +1,3 @@
+__pycache__
+rag_implementation
+.env

DockerFile

@@ -0,0 +1,16 @@
+# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+
+FROM python:3.11
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,177 @@
+'''
+Necessary Imports
+'''
+
+from fastapi import FastAPI, UploadFile, File, HTTPException,Form
+from fastapi.middleware.cors import CORSMiddleware
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from postgres import PostgresChatMessageHistory
+
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_postgres.vectorstores import PGVector
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.chains import create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+
+from typing import Dict
+from langchain_openai import ChatOpenAI
+from prompt import prompt,system_prompt
+import psycopg
+import uuid
+import os
+from custom_message import CustomMessage
+
+from dotenv import load_dotenv
+import os
+from io import BytesIO
+from pypdf import PdfReader
+from langchain.docstore.document import Document
+
+vector_store = None
+
+# LOADING ENVIRONMENT VARIABLES
+
+load_dotenv()
+
+# INSTANTIATING THE APP
+app = FastAPI()
+llm = ChatOpenAI(model="gpt-4o",
+    temperature=0.2,
+    max_tokens=None,
+    timeout=None,
+    max_retries=1)
+
+# ALLOWING CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# INITIALIZING THE EMBEDDING MODEL
+embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=1000,
+    chunk_overlap=300,
+    length_function=len,
+)
+
+
+@app.get("/")
+def greeting():
+    return {'response':'success','status code':200}
+
+# PDF UPLOAD ROUTE
+@app.post("/upload")
+async def upload_pdf(file: UploadFile = File(...), collection_name: str = Form(...)):
+    """
+    Upload and process a PDF file, storing its embeddings in the vector database.
+    """
+    if not file.filename.endswith('.pdf'):
+        raise HTTPException(status_code=400, detail="Only PDF files are allowed")
+    
+    try:
+        # Read PDF content directly into memory
+        pdf_content = await file.read()
+        pdf_file = BytesIO(pdf_content)
+        pdf_reader = PdfReader(pdf_file)
+        
+        # Extract text from PDF
+        documents = []
+        for page_num, page in enumerate(pdf_reader.pages):
+            text = page.extract_text()
+            # Create a Document object with metadata
+            doc = Document(
+                page_content=text,
+                metadata={"page": page_num + 1, "source": file.filename}
+            )
+            documents.append(doc)
+        
+        # Split documents into chunks
+        texts = text_splitter.split_documents(documents)
+        try:
+           
+        
+           global vector_store 
+           vector_store = PGVector.from_documents(
+            documents=texts,
+            embedding=embeddings,
+            connection=os.environ['CONNECTION_STRING'],
+            collection_name=collection_name,
+            use_jsonb=True,
+
+            
+        )
+           
+        
+        except Exception as e:
+            raise("Error in establishing the connection with DB: {e}")
+        
+        
+        
+        return {"message": "PDF processed successfully", "collection_name": file.filename.replace('.pdf', '')}
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    
+
+@app.post("/query")
+async def upload_pdf(query: str = Form(...),collection_name:str = Form(...),username:str = Form(...),table_name:str = Form(...)):
+    try:
+        global vector_store
+        if vector_store == None :
+            vector_store = PGVector(
+            embeddings=embeddings,
+            connection=os.environ['CONNECTION_STRING'],
+            collection_name=collection_name,
+            use_jsonb=True,
+            
+        )
+        
+        retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 3})
+        
+        
+        
+        
+        question_answer_chain = create_stuff_documents_chain(llm, prompt)
+        
+        rag_chain = create_retrieval_chain(retriever, question_answer_chain)
+        
+        response = rag_chain.invoke({"input":query})['answer']
+        
+        sync_connection = psycopg.connect(os.environ['CONNECTION_STRING'])
+        session_id = str(uuid.uuid4())
+        
+        chat_history = PostgresChatMessageHistory(
+            table_name,
+            session_id,
+            username,
+            sync_connection=sync_connection
+        )
+        
+        
+        try:
+            custom_message = CustomMessage(content=f"SYSTEM_PROMPT:{system_prompt}\n\nHUMAN_MESSAGE:{query}\n\nAI_RESPONSE:{response}")
+
+            chat_history.add_message(custom_message)
+        except Exception as e:
+            print(e)
+        print("Ended")
+        return {
+            "relevant docs":response, 
+            "session_id":session_id
+            
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+
+
+# if __name__ == "__main__":
+#     import uvicorn
+#     uvicorn.run(app, host="0.0.0.0", port=8000)
+    

custom_message.py

@@ -0,0 +1,29 @@
+from langchain_core.messages import BaseMessage
+from typing import Any, Literal, Union
+class CustomMessage(BaseMessage):
+    """Message from an AI.
+
+    AIMessage is returned from a chat model as a response to a prompt.
+
+    This message represents the output of the model and consists of both
+    the raw output as returned by the model together standardized fields
+    (e.g., tool calls, usage metadata) added by the LangChain framework.
+    """
+
+    
+
+    type: Literal["custom"] = "custom"
+    """The type of the message (used for deserialization). Defaults to "custom"."""
+
+    def __init__(
+        self, content: Union[str, list[Union[str, dict]]], **kwargs: Any
+    ) -> None:
+        """Pass in content as positional arg.
+
+        Args:
+            content: The content of the message.
+            kwargs: Additional arguments to pass to the parent class.
+        """
+        super().__init__(content=content, **kwargs)
+CustomMessage.model_rebuild()  
+        

db.py

@@ -0,0 +1,43 @@
+from langchain_community.vectorstores.pgvector import PGVector
+import os
+# try:
+#     CONNECTION_STRING = PGVector.connection_string_from_db_params(
+#     driver="psycopg2",
+#     host=os.getenv("POSTGRES_HOST"),
+#     port=int(os.getenv("POSTGRES_PORT", 5432)),
+#     database=os.getenv("POSTGRES_DB"),
+#     user=os.getenv("POSTGRES_USER"),
+#     password=os.getenv("POSTGRES_PASSWORD"),
+# )  
+#     print("Successfully established the connection")
+# except Exception as e:
+#     print("Error in establishing the connection with DB: {e}")
+print("Entered here")
+from dotenv import load_dotenv
+load_dotenv()
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
+from sqlalchemy import create_engine
+from langchain_postgres.vectorstores import PGVector
+
+from langchain_core.documents import Document
+document_1 = Document(page_content="fddsdfoo", metadata={"baz": "bar"})
+document_2 = Document(page_content="thufeed", metadata={"bar": "baz"})
+document_3 = Document(page_content="i wefsill be deleted :(")
+
+documents = [document_1, document_2, document_3]
+ids = ["1", "2", "3"]
+engine = create_engine(os.environ['CONNECTION_STRING'])
+vector_store = PGVector.from_documents(
+            documents=documents,
+            embedding=embeddings,
+            connection=os.environ['CONNECTION_STRING'],
+            collection_name="collection_name",
+            use_jsonb=True,
+
+            
+        )
+# vector_store.add_documents(documents=documents, ids=ids)
+
+
+print("Stored babe")

postgres.py

@@ -0,0 +1,379 @@
+"""Client for persisting chat message history in a Postgres database.
+
+This client provides support for both sync and async via psycopg 3.
+"""
+from __future__ import annotations
+
+import json
+import logging
+import re
+import uuid
+from typing import List, Optional, Sequence
+
+import psycopg
+from langchain_core.chat_history import BaseChatMessageHistory
+from langchain_core.messages import BaseMessage, message_to_dict, messages_from_dict
+from psycopg import sql
+
+logger = logging.getLogger(__name__)
+
+
+def _create_table_and_index(table_name: str) -> List[sql.Composed]:
+    """Make a SQL query to create a table."""
+    index_name = f"idx_{table_name}_session_id"
+    statements = [
+        sql.SQL(
+            """
+            CREATE TABLE IF NOT EXISTS {table_name} (
+                id SERIAL PRIMARY KEY,
+                username VARCHAR(255) NOT NULL, -- Add the username field
+                session_id UUID NOT NULL,
+                message JSONB NOT NULL,
+                created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+            );
+
+            """
+        ).format(table_name=sql.Identifier(table_name)),
+        sql.SQL(
+            """
+            CREATE INDEX IF NOT EXISTS {index_name} ON {table_name} (session_id);
+            """
+        ).format(
+            table_name=sql.Identifier(table_name), index_name=sql.Identifier(index_name)
+        ),
+    ]
+    return statements
+
+
+def _get_messages_query(table_name: str, last_messages: int) -> sql.Composed:
+    """Make a SQL query to get the last N messages for a given session and username."""
+    return sql.SQL(
+        "SELECT message "
+        "FROM {table_name} "
+        "WHERE session_id = %(session_id)s AND username = %(username)s "
+        "ORDER BY id DESC "
+        "LIMIT %(last_messages)s;"
+    ).format(table_name=sql.Identifier(table_name))
+
+
+def _delete_by_session_id_query(table_name: str) -> sql.Composed:
+    """Make a SQL query to delete messages for a given session and username."""
+    return sql.SQL(
+        "DELETE FROM {table_name} WHERE session_id = %(session_id)s AND username = %(username)s;"
+    ).format(table_name=sql.Identifier(table_name))
+
+
+def _delete_table_query(table_name: str) -> sql.Composed:
+    """Make a SQL query to delete a table."""
+    return sql.SQL("DROP TABLE IF EXISTS {table_name};").format(
+        table_name=sql.Identifier(table_name)
+    )
+
+
+
+def _insert_message_query(table_name: str) -> sql.Composed:
+    """Make a SQL query to insert a message with username."""
+    return sql.SQL(
+        "INSERT INTO {table_name} (username, session_id, message) VALUES (%s, %s, %s)"
+    ).format(table_name=sql.Identifier(table_name))
+
+
+
+class PostgresChatMessageHistory(BaseChatMessageHistory):
+    def __init__(
+        self,
+        table_name: str,
+        session_id: str,
+        username:str,
+        /,
+        *,
+        sync_connection: Optional[psycopg.Connection] = None,
+        async_connection: Optional[psycopg.AsyncConnection] = None,
+    ) -> None:
+        """Client for persisting chat message history in a Postgres database,
+
+        This client provides support for both sync and async via psycopg >=3.
+
+        The client can create schema in the database and provides methods to
+        add messages, get messages, and clear the chat message history.
+
+        The schema has the following columns:
+
+        - id: A serial primary key.
+        - session_id: The session ID for the chat message history.
+        - message: The JSONB message content.
+        - created_at: The timestamp of when the message was created.
+        - username: username for the user
+        Messages are retrieved for a given session_id and are sorted by
+        the id (which should be increasing monotonically), and correspond
+        to the order in which the messages were added to the history.
+
+        The "created_at" column is not returned by the interface, but
+        has been added for the schema so the information is available in the database.
+
+        A session_id can be used to separate different chat histories in the same table,
+        the session_id should be provided when initializing the client.
+
+        This chat history client takes in a psycopg connection object (either
+        Connection or AsyncConnection) and uses it to interact with the database.
+
+        This design allows to reuse the underlying connection object across
+        multiple instantiations of this class, making instantiation fast.
+
+        This chat history client is designed for prototyping applications that
+        involve chat and are based on Postgres.
+
+        As your application grows, you will likely need to extend the schema to
+        handle more complex queries. For example, a chat application
+        may involve multiple tables like a user table, a table for storing
+        chat sessions / conversations, and this table for storing chat messages
+        for a given session. The application will require access to additional
+        endpoints like deleting messages by user id, listing conversations by
+        user id or ordering them based on last message time, etc.
+
+        Feel free to adapt this implementation to suit your application's needs.
+
+        Args:
+            session_id: The session ID to use for the chat message history
+            table_name: The name of the database table to use
+            sync_connection: An existing psycopg connection instance
+            async_connection: An existing psycopg async connection instance
+
+        Usage:
+            - Use the create_tables or acreate_tables method to set up the table
+              schema in the database.
+            - Initialize the class with the appropriate session ID, table name,
+              and database connection.
+            - Add messages to the database using add_messages or aadd_messages.
+            - Retrieve messages with get_messages or aget_messages.
+            - Clear the session history with clear or aclear when needed.
+
+        Note:
+            - At least one of sync_connection or async_connection must be provided.
+
+        Examples:
+
+        .. code-block:: python
+
+            import uuid
+
+            from langchain_core.messages import SystemMessage, AIMessage, HumanMessage
+            from langchain_postgres import PostgresChatMessageHistory
+            import psycopg
+
+            # Establish a synchronous connection to the database
+            # (or use psycopg.AsyncConnection for async)
+            sync_connection = psycopg2.connect(conn_info)
+
+            # Create the table schema (only needs to be done once)
+            table_name = "chat_history"
+            PostgresChatMessageHistory.create_tables(sync_connection, table_name)
+
+            session_id = str(uuid.uuid4())
+
+            # Initialize the chat history manager
+            chat_history = PostgresChatMessageHistory(
+                table_name,
+                session_id,
+                sync_connection=sync_connection
+            )
+
+            # Add messages to the chat history
+            chat_history.add_messages([
+                SystemMessage(content="Meow"),
+                AIMessage(content="woof"),
+                HumanMessage(content="bark"),
+            ])
+
+            print(chat_history.messages)
+        """
+        if not sync_connection and not async_connection:
+            raise ValueError("Must provide sync_connection or async_connection")
+
+        self._connection = sync_connection
+        self._aconnection = async_connection
+
+        # Validate that session id is a UUID
+        try:
+            uuid.UUID(session_id)
+        except ValueError:
+            raise ValueError(
+                f"Invalid session id. Session id must be a valid UUID. Got {session_id}"
+            )
+
+        self._session_id = session_id
+        self._username = username
+        if not re.match(r"^\w+$", table_name):
+            raise ValueError(
+                "Invalid table name. Table name must contain only alphanumeric "
+                "characters and underscores."
+            )
+        self._table_name = table_name
+
+    @staticmethod
+    def create_tables(
+        connection: psycopg.Connection,
+        table_name: str,
+        /,
+    ) -> None:
+        """Create the table schema in the database and create relevant indexes."""
+        queries = _create_table_and_index(table_name)
+        logger.info("Creating schema for table %s", table_name)
+        with connection.cursor() as cursor:
+            for query in queries:
+                cursor.execute(query)
+        connection.commit()
+
+    @staticmethod
+    async def acreate_tables(
+        connection: psycopg.AsyncConnection, table_name: str, /
+    ) -> None:
+        """Create the table schema in the database and create relevant indexes."""
+        queries = _create_table_and_index(table_name)
+        logger.info("Creating schema for table %s", table_name)
+        async with connection.cursor() as cur:
+            for query in queries:
+                await cur.execute(query)
+        await connection.commit()
+
+    @staticmethod
+    def drop_table(connection: psycopg.Connection, table_name: str, /) -> None:
+        """Delete the table schema in the database.
+
+        WARNING:
+            This will delete the given table from the database including
+            all the database in the table and the schema of the table.
+
+        Args:
+            connection: The database connection.
+            table_name: The name of the table to create.
+        """
+
+        query = _delete_table_query(table_name)
+        logger.info("Dropping table %s", table_name)
+        with connection.cursor() as cursor:
+            cursor.execute(query)
+        connection.commit()
+
+    @staticmethod
+    async def adrop_table(
+        connection: psycopg.AsyncConnection, table_name: str, /
+    ) -> None:
+        """Delete the table schema in the database.
+
+        WARNING:
+            This will delete the given table from the database including
+            all the database in the table and the schema of the table.
+
+        Args:
+            connection: Async database connection.
+            table_name: The name of the table to create.
+        """
+        query = _delete_table_query(table_name)
+        logger.info("Dropping table %s", table_name)
+
+        async with connection.cursor() as acur:
+            await acur.execute(query)
+        await connection.commit()
+
+    def add_messages(self, messages: Sequence[BaseMessage]) -> None:
+        """Add messages to the chat message history."""
+        if self._connection is None:
+            raise ValueError(
+                "Please initialize the PostgresChatMessageHistory "
+                "with a sync connection or use the aadd_messages method instead."
+            )
+        # print(messages)
+        values = [
+            (self._username, self._session_id , json.dumps(message_to_dict(message)))
+            for message in messages
+        ]
+        
+        query = _insert_message_query(self._table_name)
+        # print(query.as_string(self._connection) )
+        with self._connection.cursor() as cursor:
+            cursor.executemany(query, values)
+        self._connection.commit()
+
+
+    async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:
+        """Add messages to the chat message history."""
+        if self._aconnection is None:
+            raise ValueError(
+                "Please initialize the PostgresChatMessageHistory "
+                "with an async connection or use the sync add_messages method instead."
+            )
+
+        values = [
+            (self._session_id, self._username, json.dumps(message_to_dict(message)))
+            for message in messages
+        ]
+
+        query = _insert_message_query(self._table_name)
+        async with self._aconnection.cursor() as cursor:
+            await cursor.executemany(query, values)
+        await self._aconnection.commit()
+
+    def get_messages(self, last_messages:int) -> List[BaseMessage]:
+        """Retrieve messages from the chat message history."""
+        if self._connection is None:
+            raise ValueError(
+                "Please initialize the PostgresChatMessageHistory "
+                "with a sync connection or use the async aget_messages method instead."
+            )
+
+        query = _get_messages_query(self._table_name, last_messages=last_messages)
+
+        with self._connection.cursor() as cursor:
+            cursor.execute(query, {"session_id": self._session_id})
+            items = [record[0] for record in cursor.fetchall()]
+
+        messages = messages_from_dict(items)
+        return messages
+
+    async def aget_messages(self) -> List[BaseMessage]:
+        """Retrieve messages from the chat message history."""
+        if self._aconnection is None:
+            raise ValueError(
+                "Please initialize the PostgresChatMessageHistory "
+                "with an async connection or use the sync get_messages method instead."
+            )
+
+        query = _get_messages_query(self._table_name)
+        async with self._aconnection.cursor() as cursor:
+            await cursor.execute(query, {"session_id": self._session_id,"username":self._username})
+            items = [record[0] for record in await cursor.fetchall()]
+
+        messages = messages_from_dict(items)
+        return messages
+
+    @property  # type: ignore[override]
+    def messages(self) -> List[BaseMessage]:
+        """The abstraction required a property."""
+        return self.get_messages()
+
+    def clear(self) -> None:
+        """Clear the chat message history for the GIVEN session."""
+        if self._connection is None:
+            raise ValueError(
+                "Please initialize the PostgresChatMessageHistory "
+                "with a sync connection or use the async clear method instead."
+            )
+
+        query = _delete_by_session_id_query(self._table_name)
+        with self._connection.cursor() as cursor:
+            cursor.execute(query, {"session_id": self._session_id})
+        self._connection.commit()
+
+    async def aclear(self) -> None:
+        """Clear the chat message history for the GIVEN session."""
+        if self._aconnection is None:
+            raise ValueError(
+                "Please initialize the PostgresChatMessageHistory "
+                "with an async connection or use the sync clear method instead."
+            )
+
+        query = _delete_by_session_id_query(self._table_name)
+        async with self._aconnection.cursor() as cursor:
+            await cursor.execute(query, {"session_id": self._session_id})
+        await self._aconnection.commit()

prompt.py

@@ -0,0 +1,15 @@
+from langchain_core.prompts import ChatPromptTemplate
+system_prompt = (
+    "You are a helpful assistant"
+    "If you dont know the answer , just say you dont know the answer."
+    "\n\n"
+    "Only use the following context to answer the question"
+    "{context}"
+)
+prompt = ChatPromptTemplate(
+            [
+                ("system",system_prompt),
+                ('human',"{input}"),
+               
+            ]
+        )

requirements.txt

@@ -0,0 +1,14 @@
+langchain-google-genai==2.0.1
+fastapi==0.104.1
+python-multipart==0.0.6
+langchain
+pypdf==3.17.1
+psycopg2-binary==2.9.9
+pgvector==0.2.3
+python-dotenv==1.0.0
+uvicorn==0.24.0
+langchain-postgres
+langchain-openai
+psycopg
+psycopg2
+psycopg[binary,pool]