app.py

# Import necessary libraries
import faiss
from sentence_transformers import SentenceTransformer
import numpy as np
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import re # For parsing tool calls
import gradio as gr # For the Gradio interface

# --- Configuration and Initialization ---

# File path for the knowledge base
KNOWLEDGE_FILE = 'knowledge_base.txt'

# Create a dummy knowledge file if it doesn't exist for demonstration purposes
# This ensures the app runs even without manually creating the file first in the Space
if not os.path.exists(KNOWLEDGE_FILE):
    print(f"Creating dummy knowledge file: {KNOWLEDGE_FILE}")
    with open(KNOWLEDGE_FILE, 'w', encoding='utf-8') as f:
        f.write("FAQ: How to reset my router?\n\n")
        f.write("To reset your router, locate the small reset button, usually on the back. Press and hold it for 10-15 seconds until the lights blink.\n\n")
        f.write("FAQ: What is your return policy?\n\n")
        f.write("You can return most items within 30 days of purchase if they are in original condition. Some electronics have a 15-day policy.\n\n")
        f.write("Simulated Data: Order ID 12345 is in status 'Shipped'. Items: Laptop, Mouse.\n\n")
        f.write("Simulated Data: Order ID 67890 is in status 'Processing'. Items: Keyboard.\n\n")
        f.write("Simulated Data: Order ID 11223 is in status 'Delivered'. Items: Monitor, Webcam.\n\n")
        f.write("Troubleshooting: If your device won't connect to Wi-Fi, try restarting both the device and the router.")


# --- RAG Component ---

# 1. Load knowledge and chunk (simple chunking)
def load_knowledge(filepath):
    """Loads knowledge from a text file and splits it into chunks."""
    print(f"Loading knowledge from {filepath}...")
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
        # Simple chunking by splitting on double newline
        chunks = [chunk.strip() for chunk in content.split('\n\n') if chunk.strip()]
        print(f"Loaded {len(chunks)} chunks.")
        return chunks
    except FileNotFoundError:
        print(f"Error: Knowledge file not found at {filepath}")
        return []
    except Exception as e:
        print(f"Error loading knowledge file: {e}")
        return []

# 2. Load Embedding Model
# Using a small, fast model for the pilot
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
print("Embedding model loaded: all-MiniLM-L6-v2")

# 3. Create Embeddings and FAISS Index
def create_faiss_index(chunks, model):
    """Creates FAISS index from text chunks."""
    if not chunks:
        print("No chunks to index.")
        return None, []
    print("Creating embeddings...")
    # Encode chunks to get embeddings
    embeddings = model.encode(chunks)
    dimension = embeddings.shape[1] # Dimension of the embeddings
    # Create a FAISS index (FlatL2 is simple and good for pilot)
    index = faiss.IndexFlatL2(dimension)
    # Add embeddings to the index
    index.add(np.array(embeddings).astype('float32'))
    print(f"FAISS index created with {index.ntotal} vectors.")
    return index, chunks # Return index and original chunks

# 4. Search Function in FAISS
def search_faiss(query, index, chunks, model, k=3):
    """Searches the FAISS index for the most similar chunks."""
    if index is None or not chunks:
        return ["Knowledge base not available."]
    print(f"Searching FAISS for query: '{query}' (k={k})")
    # Encode the query
    query_embedding = model.encode([query])[0]
    # Search the index
    distances, indices = index.search(np.array([query_embedding]).astype('float32'), k)
    # Retrieve the actual text chunks based on indices
    results = [chunks[i] for i in indices[0]]
    print(f"Found {len(results)} relevant chunks.")
    return results

# Prepare the knowledge base and FAISS index on startup
knowledge_chunks = load_knowledge(KNOWLEDGE_FILE)
faiss_index, indexed_chunks = create_faiss_index(knowledge_chunks, embedding_model)


# --- LLM Component (Hugging Face) ---

# Choose an LLM (gpt2 is small and easy to run on CPU)
model_name = "gpt2"

# Load Tokenizer and Model
print(f"Loading LLM: {model_name}...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Some models like gpt2 don't have a pad token, use eos token instead
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(model_name)

# Move model to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
print(f"LLM loaded: {model_name} on {device}")

# --- Tool Definitions (Simulated) ---

# Simulate a small order database
simulated_orders = {
    "12345": {"status": "Shipped", "items": ["Laptop", "Mouse"], "address": "123 Main St"},
    "67890": {"status": "Processing", "items": ["Keyboard"], "address": "456 Oak Ave"},
    "11223": {"status": "Delivered", "items": ["Monitor", "Webcam"], "address": "789 Pine Ln"},
}

def search_knowledge_base_tool(query: str) -> str:
    """
    Searches the knowledge base for relevant information based on the query.
    Args:
        query (str): The search query.
    Returns:
        str: A string containing the retrieved information chunks, separated by '---'.
             Returns a message if no information is found or if RAG is not available.
    """
    print(f"\n--- Calling Tool: search_knowledge_base with query='{query}' ---")
    if faiss_index is None or not indexed_chunks:
        return "Error: Knowledge base is not available or empty."

    results = search_faiss(query, faiss_index, indexed_chunks, embedding_model)

    if not results:
        return "No relevant information found in the knowledge base."

    formatted_results = "Retrieved Information:\n" + "\n---\n".join(results)
    print(f"--- Tool Output: {formatted_results} ---")
    return formatted_results

def get_order_status_tool(order_id: str) -> str:
    """
    Gets the current status and items of an order given its ID.
    Args:
        order_id (str): The ID of the order to look up.
    Returns:
        str: A string describing the order status and items, or an error message if not found.
    """
    print(f"\n--- Calling Tool: get_order_status with order_id='{order_id}' ---")
    order_info = simulated_orders.get(order_id)
    if order_info:
        items_list = ", ".join(order_info['items'])
        status_message = f"Order {order_id} status: {order_info['status']}. Items: {items_list}."
        print(f"--- Tool Output: {status_message} ---")
        return status_message
    else:
        error_message = f"Order {order_id} not found in our system."
        print(f"--- Tool Output: {error_message} ---")
        return error_message

def initiate_return_tool(order_id: str, item_name: str) -> str:
    """
    Initiates a return process for a specific item in an order.
    Args:
        order_id (str): The ID of the order.
        item_name (str): The name of the item to return.
    Returns:
        str: A confirmation message or an error message if the order/item is not found.
    """
    print(f"\n--- Calling Tool: initiate_return for order_id='{order_id}', item='{item_name}' ---")
    order_info = simulated_orders.get(order_id)
    if order_info:
        # Simple check if item is in the order
        if any(item.lower() == item_name.lower() for item in order_info['items']):
            # In a real system, this would interact with a return API
            confirmation_message = f"Return initiated successfully for '{item_name}' in order {order_id}. Please check your email for instructions."
            print(f"--- Tool Output: {confirmation_message} ---")
            return confirmation_message
        else:
            error_message = f"Item '{item_name}' not found in order {order_id}."
            print(f"--- Tool Output: {error_message} ---")
            return error_message
    else:
        error_message = f"Order {order_id} not found, cannot initiate return."
        print(f"--- Tool Output: {error_message} ---")
        return error_message

# Map tool names to their functions
available_tools = {
    "search_knowledge_base": search_knowledge_base_tool,
    "get_order_status": get_order_status_tool,
    "initiate_return": initiate_return_tool,
}

# --- Agentic Logic (Inspired by smolagents) ---

# System Prompt to guide the LLM
SYSTEM_PROMPT = """You are a helpful and friendly customer support agent.
Your goal is to assist the user with their queries regarding orders, products, or technical issues.
You have access to several tools to help you gather information or perform actions. Use them when necessary.

Available tools:
1.  `search_knowledge_base(query: str)`: Use this tool to find information in the knowledge base about products, FAQs, or troubleshooting steps. Use it whenever you need information you don't have or to confirm details.
2.  `get_order_status(order_id: str)`: Use this tool to check the current status and items of a customer's order. Requires an order ID.
3.  `initiate_return(order_id: str, item_name: str)`: Use this tool to start a return process for a specific item in an order. Requires the order ID and the exact item name from the order.

To use a tool, respond *only* with the following format on a new line, followed by no other text:
TOOL_USE: tool_name(arg1=value1, arg2=value2, ...)
Ensure arguments are correctly formatted (e.g., strings in quotes if necessary, though the simple parser below expects simple key=value).

If you have enough information, have completed the user's request, or cannot use a tool, respond directly to the user in natural language.
Always be polite, empathetic, and try to resolve the user's issue efficiently.
If the user asks for something unrelated to your tools or knowledge base, politely state that you can only help with support queries.
Manage the conversation flow. Ask clarifying questions if needed (e.g., ask for an order ID if the user asks about an order).
After a tool is used and its output is provided, you will receive the output. Interpret the tool output and decide the next step: either use another tool, ask a clarifying question, or provide the final answer to the user.

Conversation History:
"""

def parse_tool_call(text: str) -> tuple[str | None, dict | None]:
    """
    Parses a string to extract a tool name and arguments.
    Expected format: TOOL_USE: tool_name(arg1=value1, arg2=value2, ...)
    Returns: (tool_name, args_dict) or (None, None) if parsing fails.
    """
    tool_call_prefix = "TOOL_USE:"
    if text.strip().startswith(tool_call_prefix):
        call_string = text.strip()[len(tool_call_prefix):].strip()
        # Use regex for more robust parsing (still simplified)
        match = re.match(r"(\w+)\((.*)\)", call_string)
        if match:
            tool_name = match.group(1)
            args_string = match.group(2)
            args = {}
            # Simple parsing for key=value pairs
            if args_string:
                # Split by comma, but handle potential commas within values (not handled here for simplicity)
                arg_pairs = args_string.split(',')
                for pair in arg_pairs:
                    if '=' in pair:
                        key, value = pair.split('=', 1)
                        # Basic cleaning of value (remove leading/trailing whitespace and quotes)
                        key = key.strip()
                        value = value.strip().strip("'\"")
                        args[key] = value
            return tool_name, args
        else:
            print(f"DEBUG: Failed regex parse on tool call string: '{call_string}'")
            return None, None
    return None, None


def run_agent_step(user_input: str, conversation_history: list[str], llm_model, llm_tokenizer, tools: dict) -> tuple[str, list[str]]:
    """
    Runs one step of the agent's interaction loop.
    Determines if a tool needs to be called or generates a direct response.
    Handles multi-turn interaction including tool calls.
    """
    # Add user input to history for context
    current_history = conversation_history + [f"User: {user_input}"]

    # Prepare the full prompt for the LLM
    # Join history with newlines, ensuring system prompt is first
    prompt = SYSTEM_PROMPT + "\n" + "\n".join(current_history) + "\nAgent:"

    print(f"\n--- Sending Prompt to LLM ---\n{prompt}\n--- End Prompt ---")

    # Generate response from LLM
    inputs = llm_tokenizer(prompt, return_tensors="pt").to(device)
    max_tokens = 250 # Adjust based on desired verbosity

    output_tokens = llm_model.generate(
        **inputs,
        max_new_tokens=max_tokens,
        num_return_sequences=1,
        pad_token_id=llm_tokenizer.eos_token_id,
        # Add other generation parameters if needed (e.g., temperature, top_p)
    )

    # Decode the generated text, excluding the input prompt part
    generated_text = llm_tokenizer.decode(
        output_tokens[0][inputs.input_ids.shape[-1]:],
        skip_special_tokens=True
    ).strip()

    print(f"\n--- LLM Raw Output ---\n{generated_text}\n--- End LLM Output ---")

    # Try to parse the output as a tool call
    tool_name, tool_args = parse_tool_call(generated_text)

    if tool_name and tool_name in tools:
        # It's a tool call!
        print(f"DEBUG: Detected tool call: {tool_name} with args {tool_args}")
        try:
            # Execute the tool
            tool_output = tools[tool_name](**tool_args)

            # Add the tool call and its output to the history
            # This is the observation step for the agent
            updated_history = current_history + [f"Agent: {generated_text}", f"Tool Output: {tool_output}"]

            # Now, call the LLM *again* with the updated history including tool output.
            # This allows the agent to interpret the tool's result and formulate the final response.
            print("\n--- Re-prompting LLM with Tool Output ---")
            second_prompt = SYSTEM_PROMPT + "\n" + "\n".join(updated_history) + "\nAgent (interpreting tool output):"
            print(f"--- Second Prompt ---\n{second_prompt}\n--- End Second Prompt ---")

            second_inputs = llm_tokenizer(second_prompt, return_tensors="pt").to(device)
            second_output_tokens = llm_model.generate(
                **second_inputs,
                max_new_tokens=max_tokens, # Use same max tokens for response
                num_return_sequences=1,
                pad_token_id=llm_tokenizer.eos_token_id,
            )
            final_response = llm_tokenizer.decode(
                second_output_tokens[0][second_inputs.input_ids.shape[-1]:],
                skip_special_tokens=True
            ).strip()

            # Add the final response to the history
            final_history = updated_history + [f"Agent: {final_response}"]
            return final_response, final_history

        except TypeError as e:
             # Handle cases where args don't match tool function signature
            error_msg = f"Error executing tool '{tool_name}': Invalid arguments. {e}"
            print(f"DEBUG: {error_msg}")
            # Inform the agent about the error
            updated_history = current_history + [f"Agent: {generated_text}", f"Tool Output: Error - {error_msg}"]
            second_prompt = SYSTEM_PROMPT + "\n" + "\n".join(updated_history) + "\nAgent (interpreting tool error):"
            second_inputs = llm_tokenizer(second_prompt, return_tensors="pt").to(device)
            second_output_tokens = llm_model.generate(
                **second_inputs,
                max_new_tokens=max_tokens,
                num_return_sequences=1,
                pad_token_id=llm_tokenizer.eos_token_id,
            )
            final_response = llm_tokenizer.decode(
                second_output_tokens[0][second_inputs.input_ids.shape[-1]:],
                skip_special_tokens=True
            ).strip()
            final_history = updated_history + [f"Agent: {final_response}"]
            return final_response, final_history

        except Exception as e:
            # Handle other potential errors during tool execution
            error_msg = f"An unexpected error occurred during tool execution: {e}"
            print(f"DEBUG: {error_msg}")
            # Inform the agent about the error
            updated_history = current_history + [f"Agent: {generated_text}", f"Tool Output: Error - {error_msg}"]
            second_prompt = SYSTEM_PROMPT + "\n" + "\n".join(updated_history) + "\nAgent (interpreting tool error):"
            second_inputs = llm_tokenizer(second_prompt, return_tensors="pt").to(device)
            second_output_tokens = llm_model.generate(
                **second_inputs,
                max_new_tokens=max_tokens,
                num_return_sequences=1,
                pad_token_id=llm_tokenizer.eos_token_id,
            )
            final_response = llm_tokenizer.decode(
                second_output_tokens[0][second_inputs.input_ids.shape[-1]:],
                skip_special_tokens=True
            ).strip()
            final_history = updated_history + [f"Agent: {final_response}"]
            return final_response, final_history


    else:
        # Not a tool call, or tool not recognized. Treat as a direct response.
        print("DEBUG: No valid tool call detected. Treating as direct response.")
        final_history = current_history + [f"Agent: {generated_text}"]
        return generated_text, final_history

# --- Gradio Interface ---

# This function will be called by Gradio when the user sends a message
# It takes the user input and the current chat history (provided by Gradio)
# It returns the updated chat history
def respond(message, chat_history):
    # Gradio chat_history is a list of lists: [[user_msg, agent_msg], ...]
    # We need to convert it to our internal list of strings format for the agent
    conversation_history_for_agent = []
    for user_msg, agent_msg in chat_history:
        conversation_history_for_agent.append(f"User: {user_msg}")
        if agent_msg: # Add agent's previous response if it exists
             # Need to handle potential tool calls/outputs in previous agent messages
             # For simplicity in this pilot, we'll just add the agent's final displayed message
             # A more robust agent would track its own internal history including tool steps
             conversation_history_for_agent.append(f"Agent: {agent_msg}")


    # Run the agent step with the current message and converted history
    # Note: run_agent_step returns the agent's response and its OWN internal history list
    # We will use the response for Gradio, but the history management needs care
    # For this simple Gradio integration, we'll let run_agent_step manage history internally
    # and just use the final response. A more complex app might manage history state separately.

    # Let's modify run_agent_step slightly or manage history outside for Gradio
    # We'll pass the Gradio history and let run_agent_step return the full updated history list
    # Then convert back for Gradio.

    # --- Revised run_agent_step to handle Gradio history ---
    # The previous run_agent_step already returns the updated list history, which is good.
    # We just need to pass the Gradio history converted to list format.

    # Convert Gradio history to list of strings format for run_agent_step
    agent_internal_history = []
    for human, agent in chat_history:
        agent_internal_history.append(f"User: {human}")
        if agent is not None: # Check if agent has responded yet in a turn
             # Simple approach: just add the final agent message from the pair
            agent_internal_history.append(f"Agent: {agent}")
            # Note: This simple conversion loses the intermediate Tool Use / Tool Output steps
            # from the *previous* turns if they occurred.
            # A more advanced implementation would need to store the agent's full step history
            # separately and pass that to run_agent_step.
            # For this pilot, the current turn's tool use/output is added *within* run_agent_step
            # before the final response is generated. This is sufficient to demonstrate the logic
            # for the current turn.

    # Run the agent logic
    agent_response, updated_agent_internal_history = run_agent_step(message, agent_internal_history, model, tokenizer, available_tools)

    # Append the agent's final response to the Gradio chat history
    chat_history.append((message, agent_response))

    # Note: The updated_agent_internal_history contains the full trace for the *last* turn.
    # If you wanted to persist the full detailed history across turns for the agent's context,
    # you would need to store updated_agent_internal_history in the Gradio state.
    # For this pilot, the simple conversion above provides enough context for the LLM
    # to remember previous user/agent messages, but not the intermediate tool steps of past turns.

    return "", chat_history # Return empty string for message box, and updated history


# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Agentic RAG Pilot Demo
        This is a small-scale pilot demonstrating an agentic system with RAG and tool use
        for customer support queries, inspired by the smolagents framework.
        Ask me about order status (try Order ID 12345 or 67890) or how to reset your router.
        """
    )
    # Update chatbot with type parameter
    chatbot = gr.Chatbot(height=400, type='messages')
    # Textbox for user input
    msg = gr.Textbox(label="Your Message")
    # Button to send message
    send_button = gr.Button("Send")
    # Button to clear chat
    clear = gr.ClearButton([msg, chatbot])

    # Link the send button and textbox submission to the respond function
    # When the send button is clicked OR the user presses Enter in the textbox,
    # call the respond function with the message and current chat history.
    # The output updates the message box (clearing it) and the chatbot.
    send_button.click(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
    msg.submit(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])


# Launch the Gradio app with updated parameters
if __name__ == "__main__":
    demo.queue().launch()  # Use queue() method instead of enable_queue parameter