chain.py

"""Demo script for loading, parsing, and chunking a PDF document using LangChain."""

import os

from dotenv import load_dotenv
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders.parsers import PyPDFParser
from langchain_community.vectorstores import PGVector
from langchain_core.documents import Document
from langchain_core.documents.base import Blob
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable, chain
from langchain_openai import OpenAIEmbeddings
from langchain_openai.chat_models import ChatOpenAI
from pydantic import BaseModel

# Load environment variables (e.g., OpenAI API key)
load_dotenv()

# --- Configuration ---
PDF_PATH = "data/Agentic_AI_Resume.pdf"
COLLECTION_NAME = "agentic_ai_resume"
CONNECTION = "postgresql+psycopg2://langchain:langchain@localhost:6024/langchain"

# --- Load and parse PDF ---
blob = Blob.from_path(PDF_PATH)
parser = PyPDFParser()
documents = parser.parse(blob)

print("✅ First page content:\n", documents[0].page_content)

# --- Split the document into chunks ---
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
chunks = text_splitter.split_documents(documents)

# --- Create embeddings and vector store ---
embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")

vector_store = PGVector.from_documents(
    documents=chunks,
    embedding=embedding_model,
    connection_string=CONNECTION,
    collection_name=COLLECTION_NAME,
)

retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 10})

# --- LLM and prompt chains ---
llm = ChatOpenAI(model="gpt-4", temperature=0)

rewrite_prompt = ChatPromptTemplate.from_template(
    "Rewrite the query to be more specific: {query}"
)
rewrite_chain = rewrite_prompt | llm | (lambda msg: msg.content.strip())

answer_prompt = ChatPromptTemplate.from_template(
    "Answer the question based on the following context:\n{context}\n\nQuestion: {question}. "
)


from langchain_core.runnables import chain
from pydantic import BaseModel


# Input/output schemas
class RagInput(BaseModel):
    query: str


class RagOutput(BaseModel):
    result: str


@chain  # ✅ No parentheses
def rag_chain(input: RagInput) -> RagOutput:
    query = input["query"]
    # new_query = rewrite_chain.invoke({"query": query})
    # print(f"\n🔁 Rewritten Query: {new_query}")
    docs = retriever.invoke(query)
    context = "\n".join([doc.page_content for doc in docs])
    llm_chain = answer_prompt | llm
    response = llm_chain.invoke({"context": context, "question": query})
    return RagOutput(result=response.content).result