src/functions_langchain.py

import os
from dotenv import load_dotenv
import sentry_sdk
from sentry_sdk.integrations.serverless import serverless_function
from typing_extensions import List, TypedDict
from langchain_core.documents import Document
from langchain.chat_models import init_chat_model
import os
from pymongo import MongoClient
from pymongo.server_api import ServerApi
from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain_huggingface import HuggingFaceEmbeddings
import streamlit as st
import bs4
from langchain import hub
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langgraph.graph import START, StateGraph
from typing_extensions import List, TypedDict
from langchain_core.vectorstores import InMemoryVectorStore

load_dotenv()

mongodb_uri = os.getenv("MONGODB_URI")
sentry_dsn = os.getenv("SENTRY_DSN")
groq_api_key = os.getenv("GROQ_API_KEY") 

sentry_sdk.init(
    dsn=sentry_dsn,
    send_default_pii=True,
    traces_sample_rate=1.0,
    _experiments={
        "continuous_profiling_auto_start": True,
    },
)

client = MongoClient(mongodb_uri, server_api=ServerApi('1'))

llm = init_chat_model("llama3-8b-8192", model_provider="groq")
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
prompt = hub.pull("rlm/rag-prompt")

@serverless_function
def initialize_inmemory_vector_store() -> InMemoryVectorStore:
    return InMemoryVectorStore(embeddings)

@serverless_function
def load_and_split_documents_from_web(link) -> List[Document]:
    loader = WebBaseLoader(
        web_paths=(link,),
        bs_kwargs=dict(
            parse_only=bs4.SoupStrainer(
                class_=("post-content", "post-title", "post-header")
            )
        ),
    )
    docs = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    all_splits = text_splitter.split_documents(docs)
    return all_splits

# Define state for application
class State(TypedDict):
    question: str
    context: List[Document]
    answer: str


# Define application steps
def retrieve(state: State, vector_store: InMemoryVectorStore):
    retrieved_docs = vector_store.similarity_search(state["question"])
    return {"context": retrieved_docs}


def generate(state: State):
    docs_content = "\n\n".join(doc.page_content for doc in state["context"])
    messages = prompt.invoke({"question": state["question"], "context": docs_content})
    response = llm.invoke(messages)
    return {"answer": response.content}

@serverless_function
def graph_init(vector_store):
    # Initialize the graph
    graph_builder = StateGraph(State)

    # Add the retrieve step directly with a lambda that captures vector_store
    graph_builder.add_node("retrieve", lambda state: retrieve(state, vector_store))

    # Add the generate step
    graph_builder.add_node("generate", generate)

    # Define the sequence (using the actual callable functions, not strings)
    graph_builder.add_edge(START, "retrieve")
    graph_builder.add_edge("retrieve", "generate")

    return graph_builder.compile()