add changes

README.md

@@ -1,3 +1,6 @@
+# Note:
+This branch is only for testing. Use the chatbot in `chatbot_testing_bundler.py` to test out the setup quickly.
+
 # Invariant Gateway
 
 **LLM proxy to observe and debug what your AI agents are doing.**

chatbot_testing_bundler.py

@@ -0,0 +1,111 @@
+from httpx import Client
+from openai import OpenAI
+import os
+
+class SimpleChatbot:
+    def __init__(self):
+
+        self.invariant_authorization_token = os.getenv('INVARIANT_API_KEY')
+        if not self.invariant_authorization_token:
+            raise ValueError("INVARIANT_API_KEY is not set")
+
+        self.client = OpenAI(
+            http_client=Client(
+                headers={
+                    "Invariant-Authorization": f"Bearer {self.invariant_authorization_token}"
+                },
+            ),
+            base_url="http://localhost:8005/api/v1/gateway/temp_dataset/openai",
+        )
+        self.conversation_history = []
+    
+    def add_message(self, role, content):
+        """Add a message to the conversation history"""
+        self.conversation_history.append({"role": role, "content": content})
+    
+    def get_response(self, user_input):
+        """Get response from the AI model"""
+        try:
+            # Add user message to history
+            self.add_message("user", user_input)
+            
+            # Get response from API
+            response = self.client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=self.conversation_history,
+            )
+            
+            # Extract the AI's response
+            ai_response = response.choices[0].message.content
+            
+            # Add AI response to history
+            self.add_message("assistant", ai_response)
+            
+            return ai_response
+            
+        except Exception as e:
+            return f"Error: {str(e)}"
+    
+    def clear_history(self):
+        """Clear the conversation history"""
+        self.conversation_history = []
+        print("Conversation history cleared!")
+    
+    def show_history(self):
+        """Display the conversation history"""
+        if not self.conversation_history:
+            print("No conversation history yet.")
+            return
+        
+        print("\n--- Conversation History ---")
+        for i, message in enumerate(self.conversation_history, 1):
+            role = message["role"].capitalize()
+            content = message["content"]
+            print(f"{i}. {role}: {content}")
+        print("--- End History ---\n")
+    
+    def run(self):
+        """Main chatbot loop"""
+        print("🤖 Simple Chatbot Started!")
+        print("Commands:")
+        print("  - Type 'quit' or 'exit' to end the conversation")
+        print("  - Type 'clear' to clear conversation history")
+        print("  - Type 'history' to show conversation history")
+        print("  - Type anything else to chat with the AI")
+        print("-" * 50)
+        
+        while True:
+            try:
+                user_input = input("\nYou: ").strip()
+                
+                if not user_input:
+                    continue
+                
+                # Handle special commands
+                if user_input.lower() in ['quit', 'exit']:
+                    print("👋 Goodbye!")
+                    break
+                elif user_input.lower() == 'clear':
+                    self.clear_history()
+                    continue
+                elif user_input.lower() == 'history':
+                    self.show_history()
+                    continue
+                
+                # Get AI response
+                print("AI: ", end="", flush=True)
+                response = self.get_response(user_input)
+                print(response)
+                
+            except KeyboardInterrupt:
+                print("\n👋 Goodbye!")
+                break
+            except Exception as e:
+                print(f"Unexpected error: {e}")
+
+def main():
+    chatbot = SimpleChatbot()
+    chatbot.run()
+
+if __name__ == "__main__":
+    main()

gateway/.env

@@ -3,9 +3,7 @@ POSTGRES_PASSWORD=postgres
 POSTGRES_DB=invariantmonitor
 POSTGRES_HOST=database
 
-# This specifies the Invariant Explorer instance where the gateway will push the traces 
-# Set this to https://preview-explorer.invariantlabs.ai if you want to push to Preview.
-# If you want to push to a local instance of explorer, then specify the app-api docker container name like: 
-# http://<app-api-docker-container-name>:8000 to push to the local explorer instance.
-INVARIANT_API_URL=https://explorer.invariantlabs.ai
-GUARDRAILS_API_URL=https://explorer.invariantlabs.ai
+# This now points to the bundler instead of the explorer. Change 
+# `bundler-session-bundler-1` to the name of the bundler container, if it is different.
+INVARIANT_API_URL=http://bundler-session-bundler-1:8210
+GUARDRAILS_API_URL=http://bundler-session-bundler-1:8210

gateway/integrations/guardrails.py

@@ -5,6 +5,7 @@ import os
 import time
 from functools import wraps
 from typing import Any
+from datetime import datetime
 
 import httpx
 from fastapi import HTTPException
@@ -136,11 +137,15 @@ async def check_guardrails(
     """
     async with httpx.AsyncClient() as client:
         url = os.getenv("GUARDRAILS_API_URL", DEFAULT_API_URL).rstrip("/")
+
+        # Note: This is probably not how we should implement this. For demonstration purposes only.
+        new_message = messages[-1].copy()
+        new_message["timestamp"] = datetime.now().isoformat()
         try:
             result = await client.post(
                 f"{url}/api/v1/policy/check/batch",
                 json={
-                    "messages": messages,
+                    "messages": messages[:-1] + [new_message],
                     "policies": [g.content for g in guardrails],
                     "parameters": context.guardrails_parameters or {},
                 },

uv.lock

@@ -249,7 +249,7 @@ wheels = [
 
 [[package]]
 name = "invariant-gateway"
-version = "0.0.5.2"
+version = "0.0.8"
 source = { editable = "." }
 dependencies = [
     { name = "fastapi" },