human cleanup

vibe_bot/tests/test_llama_wrapper.py

@@ -1,71 +1,112 @@
 # Tests all functions in the llama-wrapper.py file
 # Run with: python -m pytest test_llama_wrapper.py -v
 
-from discord import message
-import pytest
 from ..llama_wrapper import (
-    chat_completion_think,
+    chat_completion,
     chat_completion_instruct,
     image_generation,
     image_edit,
-    embeddings,
+    embedding,
 )
-from dotenv import load_dotenv
-import os
-
-OPENAI_API_CHAT_ENDPOINT = os.getenv(
-    "OPENAI_API_CHAT_ENDPOINT", "https://llama-cpp.reeselink.com"
+from ..config import (
+    CHAT_ENDPOINT,
+    CHAT_MODEL,
+    CHAT_ENDPOINT_KEY,
+    IMAGE_EDIT_ENDPOINT,
+    IMAGE_EDIT_ENDPOINT_KEY,
+    IMAGE_GEN_ENDPOINT,
+    IMAGE_GEN_ENDPOINT_KEY,
+    EMBEDDING_ENDPOINT,
+    EMBEDDING_ENDPOINT_KEY,
 )
-OPENAI_API_IMAGE_ENDPOINT = os.getenv("OPENAI_API_IMAGE_ENDPOINT")
-OPENAI_API_EDIT_ENDPOINT = os.getenv("OPENAI_API_EDIT_ENDPOINT")
-OPENAI_API_EMBED_ENDPOINT = os.getenv("OPENAI_API_EMBED_ENDPOINT")
+from io import BytesIO
+import base64
+import tempfile
+from pathlib import Path
+import numpy as np
 
-# Default models
-DEFAULT_CHAT_MODEL = os.getenv("DEFAULT_CHAT_MODEL", "qwen3.5-35b-a3b")
-DEFAULT_EMBED_MODEL = os.getenv("DEFAULT_EMBED_MODEL", "text-embedding-3-small")
-DEFAULT_IMAGE_MODEL = os.getenv("DEFAULT_IMAGE_MODEL", "dall-e-3")
-DEFAULT_EDIT_MODEL = os.getenv("DEFAULT_EDIT_MODEL", "dall-e-2")
+
+TEMPDIR = Path(tempfile.mkdtemp())
 
 
 def test_chat_completion_think():
-    # This test will fail without an actual API endpoint
-    # But it's here to show the structure
-    chat_completion_think(
+    result = chat_completion(
         system_prompt="You are a helpful assistant.",
         user_prompt="Tell me about Everquest",
-        openai_url=OPENAI_API_CHAT_ENDPOINT,
-        openai_api_key="placeholder",
-        model=DEFAULT_CHAT_MODEL,
+        openai_url=CHAT_ENDPOINT,
+        openai_api_key=CHAT_ENDPOINT_KEY,
+        model=CHAT_MODEL,
         max_tokens=100,
     )
+    print(result)
 
 
 def test_chat_completion_instruct():
-    # This test will fail without an actual API endpoint
-    # But it's here to show the structure
-    chat_completion_instruct(
+    result = chat_completion_instruct(
         system_prompt="You are a helpful assistant.",
         user_prompt="Tell me about Everquest",
-        openai_url=OPENAI_API_CHAT_ENDPOINT,
-        openai_api_key="placeholder",
-        model=DEFAULT_CHAT_MODEL,
+        openai_url=CHAT_ENDPOINT,
+        openai_api_key=CHAT_ENDPOINT_KEY,
+        model=CHAT_MODEL,
         max_tokens=100,
     )
+    print(result)
 
 
 def test_image_generation():
-    # This test will fail without an actual API endpoint
-    # But it's here to show the structure
-    pass
+    result = image_generation(
+        prompt="Generate an image of a horse",
+        openai_url=IMAGE_GEN_ENDPOINT,
+        openai_api_key=IMAGE_GEN_ENDPOINT_KEY,
+    )
+    with open("image-gen.png", "wb") as f:
+        f.write(base64.b64decode(result))
 
 
 def test_image_edit():
-    # This test will fail without an actual API endpoint
-    # But it's here to show the structure
-    pass
+    with open("image-gen.png", "rb") as f:
+        image_data = BytesIO(f.read())
+    result = image_edit(
+        image=image_data,
+        prompt="Paint the words 'horse' on the horse.",
+        openai_url=IMAGE_EDIT_ENDPOINT,
+        openai_api_key=IMAGE_EDIT_ENDPOINT_KEY,
+    )
+    with open("image-edit.png", "wb") as f:
+        f.write(base64.b64decode(result))
+
+
+def _cosine_similarity(a, b):
+    """
+    Close to 1: very similar
+    Close to 0: orthogonal
+    Close to -1: opposite
+    """
+    a, b = np.array(a), np.array(b)
+    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
 
 
 def test_embeddings():
-    # This test will fail without an actual API endpoint
-    # But it's here to show the structure
-    pass
+    result1 = embedding(
+        "this is a horse",
+        openai_url=EMBEDDING_ENDPOINT,
+        openai_api_key=EMBEDDING_ENDPOINT_KEY,
+        model="qwen3-embed-4b",
+    )
+    result2 = embedding(
+        "this is a horse also",
+        openai_url=EMBEDDING_ENDPOINT,
+        openai_api_key=EMBEDDING_ENDPOINT_KEY,
+        model="qwen3-embed-4b",
+    )
+    result3 = embedding(
+        "this is a donkey",
+        openai_url=EMBEDDING_ENDPOINT,
+        openai_api_key=EMBEDDING_ENDPOINT_KEY,
+        model="qwen3-embed-4b",
+    )
+    similarity_1 = _cosine_similarity(result1, result2)
+    assert similarity_1 > 0.9
+
+    similarity_2 = _cosine_similarity(result1, result3)
+    assert similarity_2 < 0.5