wehub-resource-sync
116 lines
3.0 KiB
Python
116 lines
3.0 KiB
Python
"""
|
|
LiteLLM module tests
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import time
|
|
import unittest
|
|
import uuid
|
|
|
|
from unittest.mock import patch
|
|
|
|
from http.server import HTTPServer, BaseHTTPRequestHandler
|
|
from threading import Thread
|
|
|
|
from txtai.pipeline import LLM
|
|
|
|
|
|
class RequestHandler(BaseHTTPRequestHandler):
|
|
"""
|
|
Test HTTP handler.
|
|
"""
|
|
|
|
def do_POST(self):
|
|
"""
|
|
POST request handler.
|
|
"""
|
|
|
|
# Parse input headers
|
|
length = int(self.headers["content-length"])
|
|
data = json.loads(self.rfile.read(length))
|
|
|
|
if data.get("stream"):
|
|
# Mock streaming response
|
|
content = "application/octet-stream"
|
|
response = (
|
|
"data: "
|
|
+ json.dumps(
|
|
{
|
|
"id": str(uuid.uuid4()),
|
|
"object": "chat.completion.chunk",
|
|
"created": int(time.time() * 1000),
|
|
"model": "test",
|
|
"choices": [{"id": 0, "delta": {"content": "blue"}}],
|
|
}
|
|
)
|
|
+ "\n\ndata: [DONE]\n\n"
|
|
)
|
|
else:
|
|
# Mock standard response
|
|
content = "application/json"
|
|
response = json.dumps(
|
|
{
|
|
"id": str(uuid.uuid4()),
|
|
"object": "chat.completion",
|
|
"created": int(time.time() * 1000),
|
|
"model": "test",
|
|
"choices": [{"id": 0, "message": {"role": "assistant", "content": "blue"}, "finish_reason": "stop"}],
|
|
}
|
|
)
|
|
|
|
# Encode response as bytes
|
|
response = response.encode("utf-8")
|
|
|
|
self.send_response(200)
|
|
self.send_header("content-type", content)
|
|
self.send_header("content-length", len(response))
|
|
self.end_headers()
|
|
|
|
self.wfile.write(response)
|
|
self.wfile.flush()
|
|
|
|
|
|
class TestLiteLLM(unittest.TestCase):
|
|
"""
|
|
LiteLLM tests.
|
|
"""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
"""
|
|
Create mock http server.
|
|
"""
|
|
|
|
cls.httpd = HTTPServer(("127.0.0.1", 8000), RequestHandler)
|
|
|
|
server = Thread(target=cls.httpd.serve_forever, daemon=True)
|
|
server.start()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
"""
|
|
Shutdown mock http server.
|
|
"""
|
|
|
|
cls.httpd.shutdown()
|
|
|
|
@patch.dict(os.environ, {"OPENAI_API_KEY": "test"})
|
|
def testGeneration(self):
|
|
"""
|
|
Test generation with LiteLLM
|
|
"""
|
|
|
|
# Test model generation with LiteLLM
|
|
model = LLM("openai/gpt-4o", api_base="http://127.0.0.1:8000")
|
|
self.assertEqual(model("The sky is"), "blue")
|
|
|
|
# Test default role
|
|
self.assertEqual(model("The sky is", defaultrole="user"), "blue")
|
|
|
|
# Test streaming
|
|
self.assertEqual(" ".join(x for x in model("The sky is", stream=True)), "blue")
|
|
|
|
# Test vision
|
|
self.assertEqual(model.isvision(), False)
|