app.py

import os
import time
import streamlit as st
import nltk

from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from pinecone import Pinecone
from pinecone_text.sparse import BM25Encoder
from langchain_community.retrievers import PineconeHybridSearchRetriever
from langchain.tools.retriever import create_retriever_tool
from langgraph.prebuilt import create_react_agent

# Download the NLTK tokenizer (if not already downloaded)
nltk.download('punkt_tab')


@st.cache_resource
def init_agent(namespace1: str, top_k1: int, namespace2: str, top_k2: int):
    """
    Initialize the LangGraph agent with two Pinecone retriever tools,
    each configured with a specified namespace and top_k value.
    """
    # Retrieve API keys from environment variables
    OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
    PINE_API_KEY = os.environ.get("PINE_API_KEY")
    if not OPENAI_API_KEY or not PINE_API_KEY:
        raise ValueError("Please set the OPENAI_API_KEY and PINE_API_KEY environment variables.")

    # --- Embeddings ---
    embed = OpenAIEmbeddings(
        model='text-embedding-3-small',
        openai_api_key=OPENAI_API_KEY,
        dimensions=768
    )

    # --- Pinecone Setup ---
    index_name = 'autogen'
    pc = Pinecone(api_key=PINE_API_KEY)
    index = pc.Index(index_name)
    # Allow a moment for the index to connect
    time.sleep(1)
    index.describe_index_stats()

    # --- BM25 Sparse Encoder ---
    bm25_encoder = BM25Encoder().default()

    # --- Create first Pinecone Hybrid Retriever Tool using namespace1 and top_k1 ---
    retriever1 = PineconeHybridSearchRetriever(
        embeddings=embed,
        sparse_encoder=bm25_encoder,
        index=index,
        namespace=namespace1,
        top_k=top_k1
    )
    retriever_tool1 = create_retriever_tool(
        retriever1,
        "retrieve_context_1",
        f"Search and return information from Autogen's codebase and documentation using namespace '{namespace1}' with top_k = {top_k1}.",
    )

    # --- Create second Pinecone Hybrid Retriever Tool using namespace2 and top_k2 ---
    retriever2 = PineconeHybridSearchRetriever(
        embeddings=embed,
        sparse_encoder=bm25_encoder,
        index=index,
        namespace=namespace2,
        top_k=top_k2
    )
    retriever_tool2 = create_retriever_tool(
        retriever2,
        "retrieve_context_2",
        f"Search and return information from Autogen's codebase and documentation using namespace '{namespace2}' with top_k = {top_k2}.",
    )

    # Both retriever tools are added to the list of available tools.
    tools = [retriever_tool1, retriever_tool2]

    # --- Chat Model ---
    model = ChatOpenAI(model_name="o3-mini-2025-01-31", openai_api_key=OPENAI_API_KEY)

    # --- System Prompt ---
    prompt = """
You are an AI coding assistant specializing in the LangGraph framework. Your primary role is to help users build, code, and debug their LangGraph graphs for multi-agent AI applications. Focus on guiding users through the actual coding and implementation of LangGraph graphs rather than merely answering theoretical questions. Your responses should empower users to write, test, and optimize their LangGraph code by leveraging documentation, source code, and practical coding examples.

You have access to two powerful tools called `retrieve_context_1` and `retrieve_context_2` that function as search engines for LangGraph’s resources. These tools are essential for retrieving up-to-date code examples, API references, and implementation details to ensure that your responses reflect the latest details from LangGraph. Use them extensively to fetch relevant coding resources when necessary.

When using these retriever tools, formulate your search queries with these key terms:
- **Graph coding**: for guidance on building and structuring LangGraph graphs.
- **Nodes implementation**: for creating, managing, and customizing workflow nodes in code.
- **Multi-agent graph workflows**: for coding interactions and collaborations among agents.
- **API Code Examples**: for detailed usage of classes, methods, and functions with code snippets.
- **Graph Execution**: for instructions on running LangGraph applications and troubleshooting code execution issues.
- **Extensions and Integrations**: for integrating third-party services or custom tools in your code.
- **LangGraph Studio Coding**: for coding best practices while using the graphical interface and prototyping.
- **Core API Code**: for understanding and coding low-level components and event-driven architectures.
- **Tool Integration in Code**: for incorporating external functionalities into your LangGraph graphs.
- **Configuration via Code**: for customizing the framework’s behavior programmatically.
- **Code Migration**: for instructions on upgrading LangGraph versions in your codebase.
- **Practical Coding Examples**: for real-world code samples and demonstrations.

*Note:* Append “example” to any key term (e.g., “Nodes implementation example”) to search for illustrative coding samples. Use your expertise in software engineering and AI agent development to craft additional relevant queries as needed.

When responding to user queries:
1. **Focus on coding**: Prioritize providing code examples, step-by-step coding instructions, and debugging tips related to building LangGraph graphs.
2. **Begin** by understanding the specific coding challenge or feature the user wants to implement.
3. **Search** for relevant coding examples or API details using the retriever tools if necessary.
4. **Provide** clear, concise, and accurate code snippets, including explanations for each part of the code.
5. **Explain** technical concepts in a way that is accessible to developers who are implementing LangGraph graphs.
6. **Suggest** best practices, testing strategies, and debugging techniques for the user’s code.

**Response Format:**
- Start with a brief introduction that acknowledges the user’s coding challenge or request.
- Present the main coding solution or explanation with well-commented code snippets.
- Include any relevant code samples or API usage examples directly in your response.
- Offer additional context, tips, or advanced techniques related to coding LangGraph graphs.
- Conclude with recommendations for next steps, additional topics, or further code refinement tips.

If a user’s query is unclear or falls outside the direct scope of coding LangGraph graphs, politely ask for clarification or guide them towards more appropriate resources.

Always use the retriever tools frequently—even for queries you are confident about—since LangGraph’s coding resources are continuously updated.

Now, please help the user with their coding query for LangGraph:
    """

    # --- Create the React Agent using both tools ---
    graph = create_react_agent(model, tools=tools, messages_modifier=prompt)
    return graph


# ----------------- Sidebar: Namespace & Top_K Selection ----------------- #
st.sidebar.header("Retriever Tool Settings")

# Dropdown and slider for Retriever Tool 1
namespace_options = ["langgraph-main", "autogen"]
namespace1 = st.sidebar.selectbox("Select namespace for Retriever Tool 1:", namespace_options, index=0)
top_k1 = st.sidebar.slider("Select top_k for Retriever Tool 1:", min_value=1, max_value=4, value=1, step=1)

# Dropdown and slider for Retriever Tool 2
namespace2 = st.sidebar.selectbox("Select namespace for Retriever Tool 2:", namespace_options, index=0)
top_k2 = st.sidebar.slider("Select top_k for Retriever Tool 2:", min_value=1, max_value=4, value=1, step=1)

# Initialize the agent with the selected namespaces and top_k values.
graph = init_agent(namespace1, top_k1, namespace2, top_k2)


# ----------------- Main Chat App UI ----------------- #
st.title("LangGraph Coding Chat Assistant")

# Initialize conversation history in session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []  # Each entry is a tuple: (role, message)

def display_conversation():
    """Display the chat conversation."""
    for role, message in st.session_state.chat_history:
        if role == "user":
            st.markdown(f"**You:** {message}")
        else:
            st.markdown(f"**Assistant:** {message}")

# Display the existing conversation
display_conversation()

# --- Chat Input Form ---
with st.form("chat_form", clear_on_submit=True):
    user_input = st.text_input("Enter your message:")
    submitted = st.form_submit_button("Send")
    if submitted and user_input:
        st.session_state.chat_history.append(("user", user_input))
        st.experimental_rerun()

# --- Generate Assistant Response ---
if st.session_state.chat_history and st.session_state.chat_history[-1][0] == "user":
    inputs = {"messages": st.session_state.chat_history}

    # Placeholder for real-time streaming of the response.
    response_placeholder = st.empty()
    assistant_message = ""

    # Stream the agent's response chunk-by-chunk.
    for s in graph.stream(inputs, stream_mode="values"):
        # Extract the last message from the messages list.
        message = s["messages"][-1]
        if isinstance(message, tuple):
            role, text = message
        else:
            text = message.content
        assistant_message += text
        response_placeholder.markdown(f"**Assistant:** {assistant_message}")

    # Append the complete assistant response to the chat history.
    st.session_state.chat_history.append(("assistant", assistant_message))
    st.experimental_rerun()