Integrations

import chainlit as cl

def instrument_openai() -> None:
    """
    Automatically instrument OpenAI SDK for step tracking and observability.
    Call once during application startup to enable automatic LLM step creation.
    
    Requirements:
        - openai >= 1.0.0
        
    Features:
        - Automatic step creation for all OpenAI API calls
        - Model name, timing, and token usage capture
        - Message history tracking for chat completions
        - Support for streaming responses
        - Error handling and retry monitoring
        
    Usage:
        Call this function once at startup, then use OpenAI SDK normally.
        All calls will automatically appear as steps in the Chainlit UI.
        
    Returns:
        None
    """

import chainlit as cl
import openai

# Initialize OpenAI instrumentation (call once at startup)
cl.instrument_openai()

# Initialize OpenAI client
openai_client = openai.AsyncOpenAI(api_key="your-api-key")

@cl.on_message
async def chat_with_openai(message: cl.Message):
    """Chat with OpenAI using automatic step tracking"""
    
    # This call will automatically create a step in the UI
    response = await openai_client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": message.content}
        ],
        temperature=0.7,
        max_tokens=150
    )
    
    # Response is automatically tracked with:
    # - Model used (gpt-3.5-turbo)
    # - Token usage (prompt + completion tokens)
    # - Timing information
    # - Input messages and output response
    
    await cl.Message(response.choices[0].message.content).send()

@cl.on_message
async def streaming_openai_example(message: cl.Message):
    """Example with streaming responses"""
    
    # Streaming calls are also automatically instrumented
    stream = await openai_client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": message.content}],
        stream=True,
        temperature=0.8
    )
    
    # Stream response to UI
    msg = cl.Message("")
    await msg.send()
    
    async for chunk in stream:
        if chunk.choices[0].delta.content:
            await msg.stream_token(chunk.choices[0].delta.content)
    
    # Step automatically tracks full conversation and final token counts

@cl.on_message  
async def function_calling_example(message: cl.Message):
    """Example with OpenAI function calling"""
    
    # Function definitions
    functions = [
        {
            "name": "get_weather",
            "description": "Get current weather information",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {"type": "string", "description": "City name"}
                },
                "required": ["location"]
            }
        }
    ]
    
    # Function call - automatically tracked
    response = await openai_client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": message.content}],
        functions=functions,
        function_call="auto"
    )
    
    # Handle function call if requested
    if response.choices[0].message.function_call:
        function_name = response.choices[0].message.function_call.name
        # Function execution is tracked in the same step
        
        # Execute function (example)
        if function_name == "get_weather":
            result = "Sunny, 72°F"
            
            # Send function result back to OpenAI
            final_response = await openai_client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[
                    {"role": "user", "content": message.content},
                    response.choices[0].message,
                    {"role": "function", "name": function_name, "content": result}
                ]
            )
            
            await cl.Message(final_response.choices[0].message.content).send()
    else:
        await cl.Message(response.choices[0].message.content).send()

class LangchainCallbackHandler:
    """
    Callback handler for LangChain integration with automatic step tracking.
    
    Requirements:
        - langchain >= 0.0.198
        
    Args:
        answer_prefix_tokens: Optional[List[str]] - Tokens that mark final answer
        stream_final_answer: bool - Whether to stream final responses (default: False)
        to_ignore: List[str] - LangChain component names to ignore in tracking
        to_keep: List[str] - Components to keep even if parent is ignored
        
    Features:
        - Automatic step creation for LangChain components (chains, agents, tools)
        - LLM call tracking with generation metadata
        - Token streaming support for real-time responses
        - Tool execution and retrieval step monitoring
        - Hierarchical step organization matching LangChain execution flow
        
    Returns:
        LangchainCallbackHandler instance for use with LangChain
    """
    def __init__(
        self,
        answer_prefix_tokens: Optional[List[str]] = None,
        stream_final_answer: bool = False,
        to_ignore: List[str] = [],
        to_keep: List[str] = []
    ): ...

# Async version for async LangChain components
AsyncLangchainCallbackHandler = LangchainCallbackHandler

import chainlit as cl
from langchain.chains import LLMChain
from langchain.prompts import ChatPromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.agents import create_openai_functions_agent, AgentExecutor
from langchain.tools import Tool

@cl.on_message
async def langchain_basic_example(message: cl.Message):
    """Basic LangChain integration example"""
    
    # Create LangChain callback handler
    callback_handler = cl.AsyncLangchainCallbackHandler(
        stream_final_answer=True,
        answer_prefix_tokens=["Final", "Answer"]
    )
    
    # Create LangChain components
    llm = ChatOpenAI(
        model="gpt-3.5-turbo",
        temperature=0.7,
        streaming=True,
        callbacks=[callback_handler]
    )
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a helpful assistant."),
        ("human", "{input}")
    ])
    
    # Create and run chain - automatically creates steps
    chain = prompt | llm
    
    response = await chain.ainvoke(
        {"input": message.content},
        config={"callbacks": [callback_handler]}
    )
    
    # Response is automatically streamed to UI via callback handler
    # Steps show: Chain execution -> LLM call -> Response streaming

@cl.on_message
async def langchain_agent_example(message: cl.Message):
    """LangChain agent with tools example"""
    
    callback_handler = cl.AsyncLangchainCallbackHandler(
        stream_final_answer=True
    )
    
    # Define tools
    def calculate(expression: str) -> str:
        """Calculate mathematical expressions."""
        try:
            result = eval(expression)  # Note: Use safe eval in production
            return str(result)
        except Exception as e:
            return f"Error: {str(e)}"
    
    def search_web(query: str) -> str:
        """Search the web for information."""
        # Mock implementation
        return f"Search results for: {query}"
    
    tools = [
        Tool(
            name="Calculator",
            func=calculate,
            description="Calculate mathematical expressions"
        ),
        Tool(
            name="WebSearch", 
            func=search_web,
            description="Search the web for current information"
        )
    ]
    
    # Create agent
    llm = ChatOpenAI(
        model="gpt-3.5-turbo",
        temperature=0,
        callbacks=[callback_handler]
    )
    
    agent = create_openai_functions_agent(llm, tools, 
        ChatPromptTemplate.from_messages([
            ("system", "You are a helpful assistant with access to tools."),
            ("human", "{input}"),
            ("placeholder", "{agent_scratchpad}")
        ])
    )
    
    agent_executor = AgentExecutor(
        agent=agent, 
        tools=tools,
        callbacks=[callback_handler],
        verbose=True
    )
    
    # Execute agent - creates hierarchical steps:
    # Agent Execution -> Tool Selection -> Tool Execution -> LLM Response
    response = await agent_executor.ainvoke(
        {"input": message.content},
        config={"callbacks": [callback_handler]}
    )
    
    await cl.Message(response["output"]).send()

@cl.on_message
async def langchain_rag_example(message: cl.Message):
    """LangChain RAG (Retrieval Augmented Generation) example"""
    
    callback_handler = cl.AsyncLangchainCallbackHandler(
        stream_final_answer=True,
        to_ignore=["VectorStoreRetriever"],  # Ignore noisy retriever logs
        to_keep=["RetrievalQA"]  # But keep the main QA chain
    )
    
    # Mock retrieval setup (replace with real vector store)
    from langchain.vectorstores import FAISS
    from langchain.embeddings import OpenAIEmbeddings
    from langchain.chains import RetrievalQA
    
    # Create mock vector store
    embeddings = OpenAIEmbeddings(callbacks=[callback_handler])
    
    # Mock documents (replace with real document loading)
    docs = ["Document content 1", "Document content 2"] 
    vectorstore = FAISS.from_texts(docs, embeddings)
    
    # Create RAG chain
    llm = ChatOpenAI(
        model="gpt-3.5-turbo",
        callbacks=[callback_handler]
    )
    
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vectorstore.as_retriever(),
        callbacks=[callback_handler]
    )
    
    # Execute RAG - creates steps for:
    # Question Processing -> Document Retrieval -> Context Assembly -> LLM Generation
    response = await qa_chain.ainvoke(
        {"query": message.content},
        config={"callbacks": [callback_handler]}
    )
    
    await cl.Message(response["result"]).send()

class LlamaIndexCallbackHandler:
    """
    Callback handler for LlamaIndex integration with step tracking.
    
    Args:
        event_starts_to_ignore: List[CBEventType] - Event types to ignore at start
        event_ends_to_ignore: List[CBEventType] - Event types to ignore at end
        
    Features:
        - Automatic step creation for LlamaIndex operations
        - Query engine and retriever monitoring
        - LLM call tracking with generation data  
        - Retrieval and embedding step tracking
        - Source document display and citations
        - Index construction and update monitoring
        
    Returns:
        LlamaIndexCallbackHandler instance for use with LlamaIndex
    """
    def __init__(
        self,
        event_starts_to_ignore: List = [],
        event_ends_to_ignore: List = []
    ): ...

import chainlit as cl
from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext
from llama_index.llms import OpenAI
from llama_index.callbacks import CallbackManager

@cl.on_chat_start
async def setup_llamaindex():
    """Initialize LlamaIndex with Chainlit integration"""
    
    # Create callback handler
    callback_handler = cl.LlamaIndexCallbackHandler()
    callback_manager = CallbackManager([callback_handler])
    
    # Setup LlamaIndex components
    llm = OpenAI(
        model="gpt-3.5-turbo",
        temperature=0.7
    )
    
    service_context = ServiceContext.from_defaults(
        llm=llm,
        callback_manager=callback_manager
    )
    
    # Load documents (creates steps for document processing)
    documents = SimpleDirectoryReader("./docs").load_data()
    
    # Create index (creates steps for embedding and indexing)
    index = VectorStoreIndex.from_documents(
        documents,
        service_context=service_context
    )
    
    # Store in session
    cl.user_session.set("index", index)
    cl.user_session.set("service_context", service_context)
    
    await cl.Message("Knowledge base initialized! Ask me anything about the documents.").send()

@cl.on_message
async def llamaindex_query(message: cl.Message):
    """Query LlamaIndex with automatic step tracking"""
    
    index = cl.user_session.get("index")
    service_context = cl.user_session.get("service_context")
    
    if not index:
        await cl.Message("Please wait for the knowledge base to initialize.").send()
        return
    
    # Create query engine - automatically tracked
    query_engine = index.as_query_engine(
        service_context=service_context,
        similarity_top_k=3,
        response_mode="compact"
    )
    
    # Execute query - creates hierarchical steps:
    # Query Processing -> Document Retrieval -> Context Ranking -> LLM Generation
    response = await query_engine.aquery(message.content)
    
    # Send response with source citations
    response_text = str(response)
    
    # Extract and display source documents
    source_nodes = response.source_nodes
    if source_nodes:
        sources_text = "\n\n**Sources:**\n"
        for i, node in enumerate(source_nodes, 1):
            source_text = node.text[:200] + "..." if len(node.text) > 200 else node.text
            sources_text += f"{i}. {source_text}\n"
        
        response_text += sources_text
    
    await cl.Message(response_text).send()

@cl.on_message
async def llamaindex_chat_engine(message: cl.Message):
    """Use LlamaIndex chat engine for conversational queries"""
    
    index = cl.user_session.get("index")
    
    # Get or create chat engine
    chat_engine = cl.user_session.get("chat_engine")
    if not chat_engine:
        chat_engine = index.as_chat_engine(
            chat_mode="condense_question",
            verbose=True
        )
        cl.user_session.set("chat_engine", chat_engine)
    
    # Chat with context awareness - automatically creates steps
    response = await chat_engine.achat(message.content)
    
    await cl.Message(str(response)).send()

def instrument_mistralai() -> None:
    """
    Instrument Mistral AI SDK for automatic step tracking and observability.
    Similar to OpenAI instrumentation but for Mistral AI models.
    
    Requirements:
        - mistralai SDK
        
    Features:
        - Automatic step creation for Mistral AI API calls
        - Model performance and token usage monitoring
        - Support for Mistral's chat and completion endpoints
        - Error tracking and retry monitoring
        
    Usage:
        Call once at startup, then use Mistral AI SDK normally.
        All API calls automatically appear as steps in Chainlit UI.
        
    Returns:
        None
    """

import chainlit as cl
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage

# Initialize Mistral AI instrumentation
cl.instrument_mistralai()

# Initialize Mistral client
mistral_client = MistralClient(api_key="your-mistral-api-key")

@cl.on_message
async def chat_with_mistral(message: cl.Message):
    """Chat with Mistral AI using automatic step tracking"""
    
    # This call will automatically create a step in the UI
    response = mistral_client.chat(
        model="mistral-large-latest",
        messages=[
            ChatMessage(role="system", content="You are a helpful assistant."),
            ChatMessage(role="user", content=message.content)
        ],
        temperature=0.7,
        max_tokens=150
    )
    
    # Automatically tracked with model info, timing, and token usage
    await cl.Message(response.choices[0].message.content).send()

# Manual step creation for custom integrations
async def integrate_custom_ai_library(input_data: str) -> str:
    """Example of manual step tracking for custom AI libraries"""
    
    async with cl.Step(name="Custom AI Processing", type="llm") as step:
        step.input = {"prompt": input_data, "model": "custom-model"}
        
        # Call your custom AI library
        result = await your_ai_library.process(input_data)
        
        step.output = {"response": result, "tokens_used": 150}
        
        return result

# Step decorator for function-level tracking
@cl.step(name="Document Processing", type="tool")
async def process_document(file_path: str) -> dict:
    """Process document with automatic step tracking"""
    # Function execution is automatically wrapped in a step
    
    # Your processing logic here
    content = await extract_text(file_path)
    analysis = await analyze_content(content)
    
    return {
        "content_length": len(content),
        "analysis": analysis
    }

import chainlit as cl

# Initialize all integrations
cl.instrument_openai()
cl.instrument_mistralai()

@cl.on_message
async def multi_framework_example(message: cl.Message):
    """Example using multiple AI frameworks with unified tracking"""
    
    # LangChain for complex reasoning
    langchain_handler = cl.AsyncLangchainCallbackHandler()
    
    # OpenAI for quick responses (automatically tracked)
    quick_response = await openai_client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": f"Summarize: {message.content}"}],
        max_tokens=50
    )
    
    # LlamaIndex for document queries (with callback handler)
    index = cl.user_session.get("index")
    if index:
        llamaindex_handler = cl.LlamaIndexCallbackHandler()
        query_engine = index.as_query_engine(
            callback_manager=CallbackManager([llamaindex_handler])
        )
        doc_response = await query_engine.aquery(message.content)
    
    # Mistral AI for creative tasks (automatically tracked)
    creative_response = await mistral_client.chat(
        model="mistral-large-latest",
        messages=[{"role": "user", "content": f"Create a creative response to: {message.content}"}]
    )
    
    # All calls automatically appear as separate steps with proper attribution
    final_response = f"""
    **Quick Summary:** {quick_response.choices[0].message.content}
    
    **Document Context:** {str(doc_response) if 'doc_response' in locals() else 'No documents loaded'}
    
    **Creative Take:** {creative_response.choices[0].message.content}
    """
    
    await cl.Message(final_response).send()

from typing import List, Optional, Any, Dict
from enum import Enum

# LangChain integration types
LangchainComponent = str  # Component name for filtering

# LlamaIndex integration types  
class CBEventType(Enum):
    """LlamaIndex callback event types"""
    CHUNKING = "chunking"
    NODE_PARSING = "node_parsing" 
    EMBEDDING = "embedding"
    LLM = "llm"
    QUERY = "query"
    RETRIEVE = "retrieve"
    SYNTHESIZE = "synthesize"

# Generic integration types
ModelMetadata = Dict[str, Any]  # Model configuration and metadata
TokenUsage = Dict[str, int]     # Token usage statistics
GenerationInfo = Dict[str, Any] # Generation metadata and settings