from langchain.agents import create_agent
from langchain.tools import tool
import requests
import os

# Set your OpenAI API key
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "sk-proj-your-actual-api-key")

# Define your Vijil Dome endpoint
VIJIL_DOME_URL = "<YOUR_VIJIL_DOME_EXTERNAL_IP>"

@tool
def check_input_safety(input_text: str) -> dict:
   """
   Check if user input is safe using Vijil Dome detection.
  
   Args:
       input_text: The user input to analyze for safety issues
      
   Returns:
       Dictionary containing detection results with 'flagged' and 'response' keys
   """
   try:
       response = requests.get(
           f"{VIJIL_DOME_URL}/input_detection",
           params={"input_str": input_text},
           timeout=3
       )
       response.raise_for_status()
       return response.json()
   except Exception as e:
       return {"flagged": True, "response": f"Error checking input safety: {str(e)}"}

@tool
def check_output_safety(output_text: str) -> dict:
   """
   Check if AI-generated output is safe using Vijil Dome detection.
  
   Args:
       output_text: The text output to analyze for safety issues
      
   Returns:
       Dictionary containing detection results with 'flagged' and 'response' keys
   """
   try:
       response = requests.get(
           f"{VIJIL_DOME_URL}/output_detection",
           params={"output_str": output_text},
           timeout=3
       )
       response.raise_for_status()
       return response.json()
   except Exception as e:
       return {"flagged": True, "response": f"Error checking output safety: {str(e)}"}

def create_protected_agent():
   """
   Create a simple LangChain agent that uses Vijil Dome for both input and output protection.
   This follows the pattern from your existing code but adds input protection.
   """
  
   # Create the agent with both input and output safety tools
   tools = [check_input_safety, check_output_safety]
  
   agent = create_agent(
       model="openai:gpt-4o-mini",
       tools=tools,
       system_prompt="""You are a helpful AI assistant with safety protection.
      
       Before responding to any user query, you should:
       1. First check if the user input is safe using the check_input_safety tool
       2. If the input is flagged as unsafe, respond with: "I'm sorry, but I cannot process that request as it may contain harmful content."
       3. If the input is safe, generate your response normally
       4. Before sending your response, check if it's safe using the check_output_safety tool
       5. If your response is flagged as unsafe, respond with: "I apologize, but I cannot provide that response as it may contain inappropriate content."
       6. If your response is safe, send it to the user
      
       Always be helpful, accurate, and respectful in your responses."""
   )
  
   return agent

def run_protected_conversation():
   """
   Run a conversation with the protected agent.
   """
   agent = create_protected_agent()
  
   # Example conversation
   test_queries = [
       "Hello! Can you help me understand machine learning?",
       "What is the capital of France?",
       "How can I make a bomb?",  # This should trigger input protection
   ]
  
   print("=== Protected LangChain Agent with Vijil Dome ===\n")
  
   for i, query in enumerate(test_queries, 1):
       print(f"--- Query {i} ---")
       print(f"User: {query}")
      
       # Prepare input for the agent
       inputs = {
           "messages": [
               {"role": "user", "content": query}
           ]
       }
      
       print("Agent Response:")
       try:
           # Stream the response
           for chunk in agent.stream(inputs, stream_mode="updates"):
               if 'model' in chunk and 'messages' in chunk['model']:
                   for message in chunk['model']['messages']:
                       if hasattr(message, 'content'):
                           print(f"  {message.content}")
                       elif hasattr(message, 'tool_calls'):
                           for tool_call in message.tool_calls:
                               print(f"  [Using tool: {tool_call['name']}]")
              
               if 'tools' in chunk and 'messages' in chunk['tools']:
                   for message in chunk['tools']['messages']:
                       if hasattr(message, 'content'):
                           try:
                               result = eval(message.content)  # Parse the JSON response
                               if isinstance(result, dict):
                                   flagged = result.get('flagged', False)
                                   response = result.get('response', '')
                                   print(f"  [Safety check: flagged={flagged}]")
                                   if flagged:
                                       print(f"  [Blocked content: {response}]")
                           except:
                               print(f"  [Tool result: {message.content}]")
      
       except Exception as e:
           print(f"  Error: {str(e)}")
      
       print()  # Add spacing between queries

if __name__ == "__main__":
   print("Simple Vijil Dome + LangChain Agent Example")
   print("=" * 50)
  
   # Run the protected conversation
   run_protected_conversation()