From 31b128950b6855b1981882aa000ad2187d606d55 Mon Sep 17 00:00:00 2001 From: knielsen404 Date: Fri, 11 Jul 2025 17:49:00 +0200 Subject: [PATCH] add changes --- README.md | 3 + chatbot_testing_bundler.py | 111 +++++++++++++++++++++++++++++ gateway/.env | 10 ++- gateway/integrations/guardrails.py | 7 +- uv.lock | 2 +- 5 files changed, 125 insertions(+), 8 deletions(-) create mode 100644 chatbot_testing_bundler.py diff --git a/README.md b/README.md index 64b2a70..bba5166 100644 --- a/README.md +++ b/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.** diff --git a/chatbot_testing_bundler.py b/chatbot_testing_bundler.py new file mode 100644 index 0000000..bc6f3b8 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/gateway/.env b/gateway/.env index 093ad91..4fd8fcf 100644 --- a/gateway/.env +++ b/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://:8000 to push to the local explorer instance. -INVARIANT_API_URL=https://explorer.invariantlabs.ai -GUARDRAILS_API_URL=https://explorer.invariantlabs.ai \ No newline at end of file +# 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 \ No newline at end of file diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index add8839..425f084 100644 --- a/gateway/integrations/guardrails.py +++ b/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 {}, }, diff --git a/uv.lock b/uv.lock index 2b1418d..795a894 100644 --- a/uv.lock +++ b/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" },