From fa8cf90554a4b286037b300846b0530537786013 Mon Sep 17 00:00:00 2001 From: Christopher Tavolazzi Date: Fri, 6 Dec 2024 21:58:56 -0800 Subject: [PATCH] Update it with API Key validation --- dev/NS-bytesize/bots/bot.py | 53 +++++ dev/NS-bytesize/hubs/hub.py | 38 ++++ dev/NS-bytesize/main.py | 197 ++++++++++++++++++ dev/NS-bytesize/requirements.txt | 6 + .../tests/test_openai_key_validator.py | 79 +++++++ dev/NS-bytesize/utils/debug_env.py | 81 +++++++ dev/NS-bytesize/utils/oaic.py | 76 +++++++ dev/NS-bytesize/utils/openai_connection.py | 135 ++++++++++++ dev/NS-bytesize/utils/openai_key_validator.py | 79 +++++++ dev/NS-bytesize/utils/test_api.py | 62 ++++++ 10 files changed, 806 insertions(+) create mode 100644 dev/NS-bytesize/bots/bot.py create mode 100644 dev/NS-bytesize/hubs/hub.py create mode 100644 dev/NS-bytesize/main.py create mode 100644 dev/NS-bytesize/requirements.txt create mode 100644 dev/NS-bytesize/tests/test_openai_key_validator.py create mode 100644 dev/NS-bytesize/utils/debug_env.py create mode 100644 dev/NS-bytesize/utils/oaic.py create mode 100644 dev/NS-bytesize/utils/openai_connection.py create mode 100644 dev/NS-bytesize/utils/openai_key_validator.py create mode 100644 dev/NS-bytesize/utils/test_api.py diff --git a/dev/NS-bytesize/bots/bot.py b/dev/NS-bytesize/bots/bot.py new file mode 100644 index 00000000..26a5a075 --- /dev/null +++ b/dev/NS-bytesize/bots/bot.py @@ -0,0 +1,53 @@ +from typing import Optional, List +from uuid import uuid4 +from hubs.hub import NovaHub + +class NovaBot: + def __init__(self, hub: NovaHub): + self.hub = hub + self.session_id: Optional[str] = None + self.system_prompt = """You are a helpful AI assistant. + You provide clear, concise, and accurate responses.""" + self.message_history: List[dict] = [] + + async def initialize(self) -> str: + """Initialize the bot and return a session ID""" + self.session_id = str(uuid4()) + self.message_history = [] + return self.session_id + + async def process_message(self, message: str, model: str = "mistral") -> str: + """Process a user message and return a response""" + if not self.session_id: + raise Exception("Bot not initialized") + + # Add user message to history + self.message_history.append({ + "role": "user", + "content": message + }) + + response = await self.hub.generate_response( + prompt=message, + system=self.system_prompt, + model=model + ) + + # Add assistant response to history + self.message_history.append({ + "role": "assistant", + "content": response + }) + + return response + + async def get_history(self) -> List[dict]: + """Get the chat history""" + if not self.session_id: + raise Exception("Bot not initialized") + return self.message_history + + async def cleanup(self): + """Cleanup bot resources""" + self.session_id = None + self.message_history = [] diff --git a/dev/NS-bytesize/hubs/hub.py b/dev/NS-bytesize/hubs/hub.py new file mode 100644 index 00000000..be39750d --- /dev/null +++ b/dev/NS-bytesize/hubs/hub.py @@ -0,0 +1,38 @@ +# Standard library +from typing import Optional +import os +from pathlib import Path + +# Third party +from openai import OpenAI, AsyncOpenAI +from dotenv import load_dotenv +from utils.connection_manager import ConnectionManager +from utils.openai_key_validator import OpenAIKeyValidator + +class NovaHub: + def __init__(self, host: str = "http://localhost:11434"): + self.connection_manager = ConnectionManager() + # Get and validate API key + self.key_validator = OpenAIKeyValidator() + + # Initialize OpenAI client with validated key + self.openai_client = AsyncOpenAI(api_key=self.key_validator.key) + + # Initialize Ollama client + self.ollama_client = AsyncOpenAI( + base_url=f"{host}/v1", + api_key="ollama" + ) + + async def generate_response(self, prompt: str, model: str = "gpt-4", system: Optional[str] = None) -> str: + async with self.connection_manager.get_connection() as client: + messages = self.connection_manager.format_messages(prompt, system) + response = await client.chat.completions.create( + model=model, + messages=messages, + max_tokens=50 + ) + return response.choices[0].message.content + + async def cleanup(self): + pass # OpenAI clients don't need cleanup diff --git a/dev/NS-bytesize/main.py b/dev/NS-bytesize/main.py new file mode 100644 index 00000000..ced7bec3 --- /dev/null +++ b/dev/NS-bytesize/main.py @@ -0,0 +1,197 @@ +# Standard library +from typing import Dict, Optional, List + +# Third party +from fastapi import FastAPI, HTTPException, Depends +from fastapi.middleware.cors import CORSMiddleware +from pydantic import BaseModel +from fastapi.responses import HTMLResponse +from fastapi.staticfiles import StaticFiles +from fastapi.templating import Jinja2Templates + +# Local +from hubs.hub import NovaHub +from bots.bot import NovaBot +import logging + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) +logger = logging.getLogger(__name__) + +app = FastAPI(title="NovaSystem LITE") + +# Add CORS middleware +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], # In production, replace with specific origins + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +# Global instances +hub = NovaHub() +active_sessions: Dict[str, NovaBot] = {} + +class ChatMessage(BaseModel): + message: str + session_id: Optional[str] = None + model: Optional[str] = "gpt-4o" + +class ChatResponse(BaseModel): + response: str + session_id: str + +class ChatHistory(BaseModel): + messages: List[dict] + session_id: str + +async def get_bot(session_id: str) -> NovaBot: + """Dependency to get bot instance""" + if session_id not in active_sessions: + raise HTTPException(status_code=404, detail="Session not found") + return active_sessions[session_id] + +@app.on_event("shutdown") +async def shutdown_event(): + """Cleanup resources on shutdown""" + logger.info("Shutting down application...") + for session_id in list(active_sessions.keys()): + await active_sessions[session_id].cleanup() + await hub.cleanup() + +@app.post("/chat/", response_model=ChatResponse) +async def create_chat(): + """Create a new chat session""" + try: + bot = NovaBot(hub) + session_id = await bot.initialize() + active_sessions[session_id] = bot + logger.info(f"Created new chat session: {session_id}") + return ChatResponse(response="Chat session created", session_id=session_id) + except Exception as e: + logger.error(f"Error creating chat session: {str(e)}") + raise HTTPException(status_code=500, detail="Failed to create chat session") + +@app.post("/chat/{session_id}/message", response_model=ChatResponse) +async def send_message( + chat_message: ChatMessage, + bot: NovaBot = Depends(get_bot) +): + """Send a message in an existing chat session""" + try: + response = await bot.process_message( + chat_message.message, + model=chat_message.model + ) + return ChatResponse(response=response, session_id=bot.session_id) + except Exception as e: + logger.error(f"Error processing message: {str(e)}") + raise HTTPException(status_code=500, detail="Failed to process message") + +@app.get("/chat/{session_id}/history", response_model=ChatHistory) +async def get_history(bot: NovaBot = Depends(get_bot)): + """Get chat history for a session""" + try: + history = await bot.get_history() + return ChatHistory(messages=history, session_id=bot.session_id) + except Exception as e: + logger.error(f"Error fetching history: {str(e)}") + raise HTTPException(status_code=500, detail="Failed to fetch chat history") + +@app.post("/chat/{session_id}/end") +async def end_chat(bot: NovaBot = Depends(get_bot)): + """End a chat session""" + try: + session_id = bot.session_id + await bot.cleanup() + del active_sessions[session_id] + logger.info(f"Ended chat session: {session_id}") + return {"message": "Chat session ended"} + except Exception as e: + logger.error(f"Error ending chat session: {str(e)}") + raise HTTPException(status_code=500, detail="Failed to end chat session") + +@app.get("/", response_class=HTMLResponse) +async def root(): + return """ + + + + NovaSystem LITE Chat + + + +

NovaSystem LITE Chat

+ +
+ + + + + + + """ + +if __name__ == "__main__": + import uvicorn + uvicorn.run(app, host="0.0.0.0", port=8000) diff --git a/dev/NS-bytesize/requirements.txt b/dev/NS-bytesize/requirements.txt new file mode 100644 index 00000000..352d851b --- /dev/null +++ b/dev/NS-bytesize/requirements.txt @@ -0,0 +1,6 @@ +fastapi>=0.68.0 +uvicorn>=0.15.0 +pydantic>=1.8.0 +aiohttp>=3.8.0 +python-dotenv>=1.0.0 +openai>=1.0.0 diff --git a/dev/NS-bytesize/tests/test_openai_key_validator.py b/dev/NS-bytesize/tests/test_openai_key_validator.py new file mode 100644 index 00000000..5aa89715 --- /dev/null +++ b/dev/NS-bytesize/tests/test_openai_key_validator.py @@ -0,0 +1,79 @@ +# Standard library +import os +import unittest +from pathlib import Path +from unittest.mock import patch, MagicMock + +# Local imports +from utils.openai_key_validator import OpenAIKeyValidator + +class TestOpenAIKeyValidator(unittest.TestCase): + """Test cases for OpenAIKeyValidator.""" + + def setUp(self): + """Set up test cases.""" + # Save original environment + self.original_key = os.environ.get('OPENAI_API_KEY') + # Create a mock env path + self.mock_env_path = MagicMock(spec=Path) + self.mock_env_path.exists.return_value = False + + def tearDown(self): + """Clean up after tests.""" + # Restore original environment + if self.original_key: + os.environ['OPENAI_API_KEY'] = self.original_key + elif 'OPENAI_API_KEY' in os.environ: + del os.environ['OPENAI_API_KEY'] + + def test_valid_key(self): + """Test validator with a valid API key.""" + test_key = "sk-test123validkey456" + with patch.dict(os.environ, {'OPENAI_API_KEY': test_key}, clear=True): + validator = OpenAIKeyValidator(env_path=self.mock_env_path, search_tree=False) + self.assertTrue(validator.is_valid) + self.assertEqual(validator.key, test_key) + + def test_invalid_key_format(self): + """Test validator with invalid key format.""" + test_key = "invalid-key-format" + with patch.dict(os.environ, {'OPENAI_API_KEY': test_key}, clear=True): + validator = OpenAIKeyValidator(env_path=self.mock_env_path, search_tree=False) + self.assertFalse(validator.is_valid) + with self.assertRaises(ValueError): + _ = validator.key + + def test_placeholder_key(self): + """Test validator with placeholder key.""" + test_key = "your-actual-api-key" + with patch.dict(os.environ, {'OPENAI_API_KEY': test_key}, clear=True): + validator = OpenAIKeyValidator(env_path=self.mock_env_path, search_tree=False) + self.assertFalse(validator.is_valid) + with self.assertRaises(ValueError): + _ = validator.key + + def test_missing_key(self): + """Test validator with no API key.""" + with patch.dict(os.environ, {}, clear=True): + validator = OpenAIKeyValidator(env_path=self.mock_env_path, search_tree=False) + self.assertFalse(validator.is_valid) + with self.assertRaises(ValueError): + _ = validator.key + + def test_env_file_loading(self): + """Test .env file loading functionality.""" + test_key = "sk-test123validkey456" + mock_env_path = MagicMock(spec=Path) + mock_env_path.exists.return_value = True + + with patch('utils.openai_key_validator.load_dotenv', return_value=True) as mock_load_dotenv: + with patch.dict(os.environ, {'OPENAI_API_KEY': test_key}, clear=True): + validator = OpenAIKeyValidator(env_path=mock_env_path, search_tree=False) + + # Verify load_dotenv was called with the correct arguments + mock_load_dotenv.assert_called_once_with(mock_env_path, override=True) + self.assertTrue(validator.is_valid) + self.assertEqual(validator.key, test_key) + +if __name__ == '__main__': + unittest.main(verbosity=2) diff --git a/dev/NS-bytesize/utils/debug_env.py b/dev/NS-bytesize/utils/debug_env.py new file mode 100644 index 00000000..53e6642e --- /dev/null +++ b/dev/NS-bytesize/utils/debug_env.py @@ -0,0 +1,81 @@ +# Standard library +import os +from typing import Optional, List +from pathlib import Path + +# Third party +from dotenv import load_dotenv, find_dotenv + +def find_all_env_files() -> List[Path]: + """ + Search for all .env files in current and parent directories + """ + env_files = [] + current_dir = Path.cwd() + while current_dir.as_posix() != current_dir.root: + env_file = current_dir / '.env' + if env_file.exists(): + env_files.append(env_file) + current_dir = current_dir.parent + return env_files + +def print_api_key_status(): + """ + Debug utility to check OpenAI API key status. + Prints only the first and last 4 characters if the key exists, for security. + """ + # Clear any existing env vars + if 'OPENAI_API_KEY' in os.environ: + print("āš ļø OPENAI_API_KEY was already in environment") + del os.environ['OPENAI_API_KEY'] + + # Find and load all .env files + env_files = find_all_env_files() + + print("\n=== Environment Files Found ===") + if not env_files: + print("āš ļø No .env files found") + else: + for env_file in env_files: + print(f"\nšŸ“ {env_file}") + print(f" File exists: {env_file.exists()}") + print(f" File size: {env_file.stat().st_size} bytes") + # Read and display contents (excluding sensitive data) + with open(env_file) as f: + contents = f.readlines() + print(" Contents:") + for line in contents: + if line.strip() and not line.startswith('#'): + key = line.split('=')[0].strip() + print(f" - {key}") + if key == 'OPENAI_API_KEY': + value = line.split('=')[1].strip() + print(f" Length: {len(value)} chars") + print(f" Starts with: {value[:6]}...") + + # Load the closest .env file + if env_files: + print(f"\nLoading environment from: {env_files[0]}") + load_dotenv(env_files[0], override=True) + + api_key: Optional[str] = os.getenv('OPENAI_API_KEY') + + print("\n=== OpenAI API Key Status ===") + if not api_key: + print("āŒ No API key found in environment variables") + return + + if api_key.startswith('sk-proj-'): + print("āœ… Valid project API key format detected") + elif api_key.startswith('sk-'): + print("āœ… Valid API key format detected") + else: + print("āš ļø API key doesn't start with expected prefix") + + # Only show first 4 and last 4 characters + masked_key = f"{api_key[:6]}...{api_key[-4:]}" + print(f"āœ… API key found: {masked_key}") + print(f" Length: {len(api_key)} characters") + +if __name__ == "__main__": + print_api_key_status() diff --git a/dev/NS-bytesize/utils/oaic.py b/dev/NS-bytesize/utils/oaic.py new file mode 100644 index 00000000..40818694 --- /dev/null +++ b/dev/NS-bytesize/utils/oaic.py @@ -0,0 +1,76 @@ +# Standard library +import os +from pathlib import Path +from typing import Optional + +# Third party +from dotenv import load_dotenv + +class OpenAIKeyManager: + """ + Simple manager for OpenAI API key handling. + Validates and provides access to the OpenAI API key. + """ + def __init__(self): + self._api_key: Optional[str] = None + self._load_api_key() + + def _load_api_key(self) -> None: + """Load and validate the OpenAI API key from environment.""" + # Try to find .env file, starting from current directory and moving up + current_dir = Path.cwd() + while current_dir.as_posix() != current_dir.root: + env_file = current_dir / '.env' + if env_file.exists(): + print(f"Found .env file at: {env_file}") + load_dotenv(env_file, override=True) + break + current_dir = current_dir.parent + + # Get API key from environment + self._api_key = os.getenv('OPENAI_API_KEY') + + @property + def is_valid(self) -> bool: + """Check if we have a valid API key.""" + if not self._api_key: + return False + if self._api_key in ["your-actual-api-key", "your_api_key_here"]: + return False + if not self._api_key.startswith('sk-'): + return False + return True + + @property + def key(self) -> str: + """Get the API key, raising an error if invalid.""" + if not self.is_valid: + raise ValueError("No valid OpenAI API key found in environment") + return self._api_key + + def status(self) -> None: + """Print the current status of the API key.""" + print("\n=== OpenAI API Key Status ===") + + if not self._api_key: + print("āŒ No API key found in environment") + return + + if not self.is_valid: + print("āŒ Invalid API key format or placeholder value") + return + + # Only show first 4 and last 4 characters + masked_key = f"{self._api_key[:4]}...{self._api_key[-4:]}" + print(f"āœ… Valid API key found: {masked_key}") + print(f" Length: {len(self._api_key)} characters") + +# Example usage +if __name__ == "__main__": + key_manager = OpenAIKeyManager() + key_manager.status() + + if key_manager.is_valid: + print("\nAPI key is ready to use!") + else: + print("\nPlease set a valid OpenAI API key in your .env file") \ No newline at end of file diff --git a/dev/NS-bytesize/utils/openai_connection.py b/dev/NS-bytesize/utils/openai_connection.py new file mode 100644 index 00000000..6f1b2c21 --- /dev/null +++ b/dev/NS-bytesize/utils/openai_connection.py @@ -0,0 +1,135 @@ +# Standard library +import os +from pathlib import Path +from typing import Optional, Dict, Any + +# Third party +from openai import OpenAI, AsyncOpenAI +from dotenv import load_dotenv + +class OpenAIConnection: + """ + Manages OpenAI API connections and provides utility methods for API interactions. + """ + def __init__(self): + self._sync_client: Optional[OpenAI] = None + self._async_client: Optional[AsyncOpenAI] = None + self._api_key: Optional[str] = None + self._initialize() + + def _initialize(self) -> None: + """Initialize the connection by loading environment variables and setting up clients.""" + # Find and load .env file + env_file = Path(__file__).parent.parent.parent.parent / '.env' + if env_file.exists(): + load_dotenv(env_file, override=True) + + # Get API key + self._api_key = os.getenv('OPENAI_API_KEY') + if not self._api_key: + raise ValueError("OpenAI API key not found in environment variables") + + @property + def sync_client(self) -> OpenAI: + """Get or create synchronous OpenAI client.""" + if not self._sync_client: + self._sync_client = OpenAI(api_key=self._api_key) + return self._sync_client + + @property + def async_client(self) -> AsyncOpenAI: + """Get or create asynchronous OpenAI client.""" + if not self._async_client: + self._async_client = AsyncOpenAI(api_key=self._api_key) + return self._async_client + + def create_chat_message(self, role: str, content: str) -> Dict[str, Any]: + """ + Create a properly formatted chat message. + + Args: + role: The role (e.g., "system", "user", "assistant") + content: The message content + + Returns: + Dict containing the formatted message + """ + return { + "role": role, + "content": [{ + "type": "text", + "text": content + }] + } + + def format_messages(self, prompt: str, system_prompt: Optional[str] = None) -> list: + """ + Format messages for the chat completion API. + + Args: + prompt: The user's prompt + system_prompt: Optional system prompt + + Returns: + List of formatted messages + """ + messages = [] + if system_prompt: + messages.append(self.create_chat_message("system", system_prompt)) + messages.append(self.create_chat_message("user", prompt)) + return messages + + async def test_connection(self) -> bool: + """ + Test the OpenAI API connection. + + Returns: + bool: True if connection is successful + """ + try: + response = await self.async_client.chat.completions.create( + model="gpt-3.5-turbo", + messages=[self.create_chat_message("user", "Test connection")], + max_tokens=5 + ) + return bool(response.choices[0].message.content) + except Exception as e: + print(f"Connection test failed: {str(e)}") + return False + +# Example usage +if __name__ == "__main__": + import asyncio + + async def main(): + try: + # Initialize connection + connection = OpenAIConnection() + + # Test connection + print("\n=== Testing OpenAI Connection ===") + is_connected = await connection.test_connection() + print(f"Connection status: {'āœ… Connected' if is_connected else 'āŒ Failed'}") + + if is_connected: + # Try a simple completion + messages = connection.format_messages( + prompt="Say 'Hello, Nova!' in a cheerful way", + system_prompt="You are a helpful and enthusiastic assistant" + ) + + response = await connection.async_client.chat.completions.create( + model="gpt-4", + messages=messages, + max_tokens=50 + ) + + print("\n=== Test Response ===") + print(response.choices[0].message.content) + + except Exception as e: + print(f"\nāŒ Error: {str(e)}") + import traceback + traceback.print_exc() + + asyncio.run(main()) \ No newline at end of file diff --git a/dev/NS-bytesize/utils/openai_key_validator.py b/dev/NS-bytesize/utils/openai_key_validator.py new file mode 100644 index 00000000..5acae4ee --- /dev/null +++ b/dev/NS-bytesize/utils/openai_key_validator.py @@ -0,0 +1,79 @@ +# Standard library +import os +from pathlib import Path +from typing import Optional + +# Third party +from dotenv import load_dotenv + +class OpenAIKeyValidator: + """ + Validates and manages OpenAI API key access. + Searches for .env files, validates API key format, and provides secure access. + """ + def __init__(self, env_path: Optional[Path] = None, search_tree: bool = True): + self._api_key: Optional[str] = None + self._env_path = env_path + self._search_tree = search_tree + self._load_api_key() + + def _load_api_key(self) -> None: + """Load and validate the OpenAI API key from environment.""" + # Try to load from env_path first if provided + if self._env_path and self._env_path.exists(): + load_dotenv(self._env_path, override=True) + self._api_key = os.getenv('OPENAI_API_KEY') + return + + # Search directory tree if enabled + if self._search_tree: + current_dir = Path.cwd() + while current_dir.as_posix() != current_dir.root: + env_file = current_dir / '.env' + if env_file.exists(): + print(f"Found .env file at: {env_file}") + load_dotenv(env_file, override=True) + break + current_dir = current_dir.parent + + # Get API key from environment + self._api_key = os.getenv('OPENAI_API_KEY') + + @property + def is_valid(self) -> bool: + """Check if we have a valid API key.""" + if not self._api_key: + return False + if self._api_key in ["your-actual-api-key", "your_api_key_here"]: + return False + if not self._api_key.startswith('sk-'): + return False + return True + + @property + def key(self) -> str: + """Get the API key, raising an error if invalid.""" + if not self.is_valid: + raise ValueError("No valid OpenAI API key found in environment") + return self._api_key + + def status(self) -> None: + """Print the current status of the API key.""" + print("\n=== OpenAI API Key Status ===") + + if not self._api_key: + print("āŒ No API key found in environment") + return + + if not self.is_valid: + print("āŒ Invalid API key format or placeholder value") + return + + # Only show first 4 and last 4 characters + masked_key = f"{self._api_key[:4]}...{self._api_key[-4:]}" + print(f"āœ… Valid API key found: {masked_key}") + print(f" Length: {len(self._api_key)} characters") + +if __name__ == "__main__": + validator = OpenAIKeyValidator() + validator.status() diff --git a/dev/NS-bytesize/utils/test_api.py b/dev/NS-bytesize/utils/test_api.py new file mode 100644 index 00000000..de60c577 --- /dev/null +++ b/dev/NS-bytesize/utils/test_api.py @@ -0,0 +1,62 @@ +# Standard library +import os +from pathlib import Path + +# Third party +from openai import OpenAI +from dotenv import load_dotenv # Changed from python-dotenv to dotenv + +def test_openai_connection(): + """ + Test the OpenAI API connection with a simple completion request. + """ + print("\n=== Testing OpenAI API Connection ===") + + # Load environment + env_file = Path(__file__).parent.parent.parent.parent / '.env' + if env_file.exists(): + print(f"Loading environment from: {env_file}") + load_dotenv(env_file, override=True) + + api_key = os.getenv('OPENAI_API_KEY') + if not api_key: + print("āŒ No API key found in environment") + return + + try: + client = OpenAI(api_key=api_key) + response = client.chat.completions.create( + model="gpt-4", + messages=[ + { + "role": "system", + "content": [{ + "type": "text", + "text": """ + You are a helpful assistant that answers programming questions + in the style of a southern belle from the southeast United States. + """ + }] + }, + { + "role": "user", + "content": [{ + "type": "text", + "text": "Say 'Hello, Nova!' in your most charming southern accent." + }] + } + ], + max_tokens=50 + ) + print("\nāœ… API Connection Successful!") + print(f"Response: {response.choices[0].message.content}") + + except Exception as e: + print(f"\nāŒ API Test Failed:") + print(f"Error: {str(e)}") + print("\nFull error details:") + import traceback + traceback.print_exc() + +if __name__ == "__main__": + test_openai_connection() \ No newline at end of file