app.py

from pstuts_rag.configuration import Configuration
from pstuts_rag.datastore import fill_the_db
from pstuts_rag.graph import build_the_graph
from pstuts_rag.state import PsTutsTeamState
import requests
import asyncio
import json
import os
import getpass
from typing import List, Tuple
import re

import chainlit as cl
from dotenv import load_dotenv
from langchain_core.documents import Document
from langchain_core.language_models import BaseChatModel
from langchain_core.runnables import Runnable
from langchain_openai import ChatOpenAI
from langchain_core.embeddings import Embeddings
from langchain_huggingface import HuggingFaceEmbeddings


from langchain_core.messages import HumanMessage, BaseMessage
import langgraph.graph


import pstuts_rag.datastore
import pstuts_rag.rag


import nest_asyncio
from uuid import uuid4

import logging

logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("langchain").setLevel(logging.WARNING)

# Apply nested asyncio to enable nested event loops
nest_asyncio.apply()

# Generate a unique ID for this application instance
unique_id = uuid4().hex[0:8]

VIDEOARCHIVE = "VideoArchiveSearch"
ADOBEHELP = "AdobeHelp"


def set_api_key_if_not_present(key_name, prompt_message=""):
    """
    Sets an API key in the environment if it's not already present.

    Args:
        key_name: Name of the environment variable to set
        prompt_message: Custom prompt message for getpass (defaults to key_name)
    """
    if len(prompt_message) == 0:
        prompt_message = key_name
    if key_name not in os.environ or not os.environ[key_name]:
        os.environ[key_name] = getpass.getpass(prompt_message)


class ApplicationState:
    """
    Maintains the state of the application and its components.

    Attributes:
        embeddings: Embeddings model for vector operations
        docs: List of loaded documents
        qdrant_client: Client for Qdrant vector database
        vector_store: Vector store for document retrieval
        datastore_manager: Manager for data storage and retrieval
        rag_factory: Factory for creating RAG chains
        llm: Language model instance
        rag_chain: Retrieval-augmented generation chain
        ai_graph: Compiled AI agent graph
        ai_graph_sketch: State graph for AI agent orchestration
        tasks: List of asyncio tasks
        hasLoaded: Event to track when loading is complete
        pointsLoaded: Number of data points loaded into the database
    """

    embeddings: Embeddings = None
    docs: List[Document] = []
    qdrant_client = None
    vector_store = None
    datastore_manager = None
    rag = None
    llm: BaseChatModel = None
    rag_chain: Runnable = None

    ai_graph: Runnable = None
    ai_graph_sketch = None

    tasks: List[asyncio.Task] = []

    hasLoaded: asyncio.Event = asyncio.Event()
    pointsLoaded: int = 0

    def __init__(self) -> None:
        """
        Initialize the application state and set up environment variables.
        """
        load_dotenv()
        set_api_key_if_not_present("OPENAI_API_KEY")
        set_api_key_if_not_present("TAVILY_API_KEY")
        os.environ["LANGCHAIN_TRACING_V2"] = "true"
        os.environ["LANGCHAIN_PROJECT"] = (
            f"AIE - MBUDISIC - HF - CERT - {unique_id}"
        )
        set_api_key_if_not_present("LANGCHAIN_API_KEY")


# Initialize global application state
app_state = ApplicationState()
params = Configuration()
ai_state = PsTutsTeamState(
    messages=[],
    team_members=[VIDEOARCHIVE, ADOBEHELP],
    next="START",
)


async def initialize():

    await fill_the_db(app_state)
    app_state.ai_graph, app_state.ai_graph_sketch = await build_the_graph(
        app_state
    )


def enter_chain(message: str):
    """
    Entry point for the agent graph chain.

    Transforms a user message into the state format expected by the agent graph.

    Args:
        message: User's input message

    Returns:
        Dictionary with the message and team members information
    """
    results = {
        "messages": [HumanMessage(content=message)],
        "team_members": [VIDEOARCHIVE, ADOBEHELP],
    }
    return results


@cl.on_chat_start
async def on_chat_start():
    """
    Initializes the application when a new chat session starts.

    Sets up the language model, vector database components, and spawns tasks
    for database population and graph building.
    """
    app_state.llm = ChatOpenAI(model=params.tool_calling_model, temperature=0)
    # Use LangChain's built-in HuggingFaceEmbeddings wrapper
    app_state.embeddings = HuggingFaceEmbeddings(
        model_name=params.embedding_model
    )

    app_state.rag = pstuts_rag.rag.RAGChainInstance(
        name="deployed",
        qdrant_client=app_state.qdrant_client,
        llm=app_state.llm,
        embeddings=app_state.embeddings,
    )

    app_state.tasks.append(asyncio.create_task(initialize()))


def process_response(
    response_message: BaseMessage,
) -> Tuple[str, List[cl.Message]]:
    """
    Processes a response from the AI agents.

    Extracts the main text and video references from the response,
    and creates message elements for displaying video content.

    Args:
        response: Response object from the AI agent

    Returns:
        Tuple containing the text response and a list of message elements with video references
    """
    streamed_text = f"[_from: {response_message.name}_]\n"
    msg_references = []

    if response_message.name == VIDEOARCHIVE:
        text, references = pstuts_rag.rag.RAGChainFactory.unpack_references(
            str(response_message.content)
        )
        streamed_text += text

        if len(references) > 0:
            references = json.loads(references)
            print(references)

            for ref in references:
                msg_references.append(
                    cl.Message(
                        content=(
                            f"Watch {ref['title']} from timestamp "
                            f"{round(ref['start'] // 60)}m:{round(ref['start'] % 60)}s"
                        ),
                        elements=[
                            cl.Video(
                                name=ref["title"],
                                url=f"{ref['source']}#t={ref['start']}",
                                display="side",
                            )
                        ],
                    )
                )
    else:
        streamed_text += str(response_message.content)

        # Find all URLs in the content
        urls = re.findall(
            r"https?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+[/\w\.-]*(?:\?[/\w\.-=&%]*)?",
            str(response_message.content),
        )
        print(urls)
        links = []
        # Create a list of unique URLs
        for idx, u in enumerate(list(set(urls))):

            url = "https://api.microlink.io"
            params = {
                "url": u,
                "screenshot": True,
            }

            payload = requests.get(url, params)

            if payload:
                print(f"Successful screenshot\n{payload.json()}")
                links.append(
                    cl.Image(
                        name=f"Website {idx} Preview: {u}",
                        display="side",  # Show in the sidebar
                        url=payload.json()["data"]["screenshot"]["url"],
                    )
                )

        print(links)
        msg_references.append(
            cl.Message(
                content="\n".join([l.url for l in links]), elements=links
            )
        )

    return streamed_text, msg_references


@cl.on_message
async def main(user_cl_message: cl.Message):
    """
    Processes incoming user messages and sends responses.

    Streams the AI agent's response, processes it to extract text and video references,
    and sends the content back to the user's chat interface.

    Args:
        message: User's input message
    """
    for s in app_state.ai_graph.stream(
        user_cl_message.content, {"recursion_limit": 20}
    ):
        if "__end__" not in s and "supervisor" not in s.keys():
            for [node_type, node_response] in s.items():
                print(f"Processing {node_type} messages")
                for node_message in node_response["messages"]:
                    print(f"Message {node_message}")
                    msg = cl.Message(content="")
                    text, references = process_response(node_message)
                    for token in [char for char in text]:
                        await msg.stream_token(token)
                    await msg.send()
                    for m in references:
                        await m.send()


if __name__ == "__main__":
    main()