diff --git a/backend/lib/types.py b/backend/lib/types.py index 7e46583..a94f73b 100644 --- a/backend/lib/types.py +++ b/backend/lib/types.py @@ -1,5 +1,5 @@ from pydantic import BaseModel, Field -from typing import List, Optional, Dict, Union +from typing import List, Optional, Union # Add these models after your existing Pydantic models class TranscriptionRequest(BaseModel): @@ -20,7 +20,7 @@ class ImageUrl(BaseModel): class MessageContent(BaseModel): type: str # "text" or "image_url" text: Optional[str] = None - image_url: Optional[ImageUrl] = None + image_url: Optional[ImageUrl] = None # chat completion request and response models class UsageDetails(BaseModel): @@ -30,14 +30,14 @@ class UsageDetails(BaseModel): class Usage(BaseModel): prompt_tokens: int - completion_tokens: int + completion_tokens: Optional[int] = None total_tokens: int prompt_tokens_details: Optional[UsageDetails] = None completion_tokens_details: Optional[UsageDetails] = None class Message(BaseModel): role: str - content: Union[str, List[MessageContent]] # Can be string or list for vision + content: Union[str,List[MessageContent]] # Can be string or list for vision name: Optional[str] = None class ChatCompletionRequest(BaseModel): @@ -69,7 +69,7 @@ class ChatCompletionResponse(BaseModel): # embeddings request and response models class EmbeddingInput(BaseModel): model: str - input: List[str] + input: Union[str,List[str]] user: Optional[str] = None class EmbeddingData(BaseModel): @@ -88,4 +88,4 @@ class SpeechRequest(BaseModel): input: str voice: str = "alloy" # alloy, echo, fable, onyx, nova, shimmer response_format: Optional[str] = "mp3" # mp3, opus, aac, flac, wav, pcm - speed: Optional[float] = Field(1.0, ge=0.25, le=4.0) # Speed between 0.25 and 4.0 \ No newline at end of file + speed: Optional[float] = Field(1.0, ge=0.25, le=4.0) # Speed between 0.25 and 4.0 diff --git a/backend/main.py b/backend/main.py index 6bfca66..ba0ea99 100644 --- a/backend/main.py +++ b/backend/main.py @@ -128,7 +128,7 @@ async def chat_completions(request: ChatCompletionRequest, user_key = Depends(ve if content_item.image_url.url.startswith(("http://", "https://")): content_item.image_url.url = await fetch_image_as_base64(content_item.image_url.url) - request_data = request.dict(by_alias=True) + request_data = request.model_dump(by_alias=True) # Convert to OpenAI format for vision messages if has_images: @@ -160,7 +160,11 @@ async def chat_completions(request: ChatCompletionRequest, user_key = Depends(ve # Keep OpenAI-compatible parameters only allowed_params = ["model", "messages", "stream", "max_tokens", "temperature", "top_p", "n", "stop", "presence_penalty", "frequency_penalty", "user"] request_data = {k: v for k, v in request_data.items() if k in allowed_params and v is not None} - + + # Some models don't allow both temperature and top_p + if 'temperature' in request_data and 'top_p' in request_data: + # Remove top_p, keep temperature (or vice versa based on your preference) + request_data.pop('top_p') # Don't truncate messages with images if model_config['params'].get('max_input_tokens') and not has_images: # Only truncate text-only messages @@ -343,7 +347,7 @@ async def event_generator(): @app.post("/v1/embeddings") async def create_embedding(request: EmbeddingInput, user_key = Depends(verify_token)): model_config = get_model_config(request.model, user_key) - request_data = request.dict() + request_data = request.model_dump() request_data["model"] = model_config['params']['model'] # Maps "devstral" to "devstral:24b" @@ -599,4 +603,4 @@ async def list_models(user_key = Depends(verify_token)): if __name__ == "__main__": import uvicorn - uvicorn.run(app, host="0.0.0.0", port=8000) \ No newline at end of file + uvicorn.run(app, host="0.0.0.0", port=8000) diff --git a/backend/requirements.txt b/backend/requirements.txt index f8c5084..8956ad4 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -7,4 +7,6 @@ codecarbon prometheus-client python-multipart tiktoken -requests \ No newline at end of file +requests +pytest +httpx diff --git a/backend/test_main.py b/backend/test_main.py new file mode 100644 index 0000000..74fd65b --- /dev/null +++ b/backend/test_main.py @@ -0,0 +1,122 @@ +from fastapi.testclient import TestClient +from main import app +from lib.types import ChatCompletionResponse, EmbeddingResponse +import os + + +token = os.environ["TOKEN"] +print("TOKEN", token ) +client = TestClient(app) + +b64_image = "data:image/gif;base64,R0lGODlhPQBEAPeoAJosM//AwO/AwHVYZ/z595kzAP/s7P+goOXMv8+fhw/v739/f+8PD98fH/8mJl+fn/9ZWb8/PzWlwv///6wWGbImAPgTEMImIN9gUFCEm/gDALULDN8PAD6atYdCTX9gUNKlj8wZAKUsAOzZz+UMAOsJAP/Z2ccMDA8PD/95eX5NWvsJCOVNQPtfX/8zM8+QePLl38MGBr8JCP+zs9myn/8GBqwpAP/GxgwJCPny78lzYLgjAJ8vAP9fX/+MjMUcAN8zM/9wcM8ZGcATEL+QePdZWf/29uc/P9cmJu9MTDImIN+/r7+/vz8/P8VNQGNugV8AAF9fX8swMNgTAFlDOICAgPNSUnNWSMQ5MBAQEJE3QPIGAM9AQMqGcG9vb6MhJsEdGM8vLx8fH98AANIWAMuQeL8fABkTEPPQ0OM5OSYdGFl5jo+Pj/+pqcsTE78wMFNGQLYmID4dGPvd3UBAQJmTkP+8vH9QUK+vr8ZWSHpzcJMmILdwcLOGcHRQUHxwcK9PT9DQ0O/v70w5MLypoG8wKOuwsP/g4P/Q0IcwKEswKMl8aJ9fX2xjdOtGRs/Pz+Dg4GImIP8gIH0sKEAwKKmTiKZ8aB/f39Wsl+LFt8dgUE9PT5x5aHBwcP+AgP+WltdgYMyZfyywz78AAAAAAAD///8AAP9mZv///wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACH5BAEAAKgALAAAAAA9AEQAAAj/AFEJHEiwoMGDCBMqXMiwocAbBww4nEhxoYkUpzJGrMixogkfGUNqlNixJEIDB0SqHGmyJSojM1bKZOmyop0gM3Oe2liTISKMOoPy7GnwY9CjIYcSRYm0aVKSLmE6nfq05QycVLPuhDrxBlCtYJUqNAq2bNWEBj6ZXRuyxZyDRtqwnXvkhACDV+euTeJm1Ki7A73qNWtFiF+/gA95Gly2CJLDhwEHMOUAAuOpLYDEgBxZ4GRTlC1fDnpkM+fOqD6DDj1aZpITp0dtGCDhr+fVuCu3zlg49ijaokTZTo27uG7Gjn2P+hI8+PDPERoUB318bWbfAJ5sUNFcuGRTYUqV/3ogfXp1rWlMc6awJjiAAd2fm4ogXjz56aypOoIde4OE5u/F9x199dlXnnGiHZWEYbGpsAEA3QXYnHwEFliKAgswgJ8LPeiUXGwedCAKABACCN+EA1pYIIYaFlcDhytd51sGAJbo3onOpajiihlO92KHGaUXGwWjUBChjSPiWJuOO/LYIm4v1tXfE6J4gCSJEZ7YgRYUNrkji9P55sF/ogxw5ZkSqIDaZBV6aSGYq/lGZplndkckZ98xoICbTcIJGQAZcNmdmUc210hs35nCyJ58fgmIKX5RQGOZowxaZwYA+JaoKQwswGijBV4C6SiTUmpphMspJx9unX4KaimjDv9aaXOEBteBqmuuxgEHoLX6Kqx+yXqqBANsgCtit4FWQAEkrNbpq7HSOmtwag5w57GrmlJBASEU18ADjUYb3ADTinIttsgSB1oJFfA63bduimuqKB1keqwUhoCSK374wbujvOSu4QG6UvxBRydcpKsav++Ca6G8A6Pr1x2kVMyHwsVxUALDq/krnrhPSOzXG1lUTIoffqGR7Goi2MAxbv6O2kEG56I7CSlRsEFKFVyovDJoIRTg7sugNRDGqCJzJgcKE0ywc0ELm6KBCCJo8DIPFeCWNGcyqNFE06ToAfV0HBRgxsvLThHn1oddQMrXj5DyAQgjEHSAJMWZwS3HPxT/QMbabI/iBCliMLEJKX2EEkomBAUCxRi42VDADxyTYDVogV+wSChqmKxEKCDAYFDFj4OmwbY7bDGdBhtrnTQYOigeChUmc1K3QTnAUfEgGFgAWt88hKA6aCRIXhxnQ1yg3BCayK44EWdkUQcBByEQChFXfCB776aQsG0BIlQgQgE8qO26X1h8cEUep8ngRBnOy74E9QgRgEAC8SvOfQkh7FDBDmS43PmGoIiKUUEGkMEC/PJHgxw0xH74yx/3XnaYRJgMB8obxQW6kL9QYEJ0FIFgByfIL7/IQAlvQwEpnAC7DtLNJCKUoO/w45c44GwCXiAFB/OXAATQryUxdN4LfFiwgjCNYg+kYMIEFkCKDs6PKAIJouyGWMS1FSKJOMRB/BoIxYJIUXFUxNwoIkEKPAgCBZSQHQ1A2EWDfDEUVLyADj5AChSIQW6gu10bE/JG2VnCZGfo4R4d0sdQoBAHhPjhIB94v/wRoRKQWGRHgrhGSQJxCS+0pCZbEhAAOw==" +url_image = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg" +chatRequest = { + "model": "gemma3", # Replace with actual model name (e.g., "gpt-4") + "messages": [ + { + "role": "user", # or "system", "assistant" + "content": "Hello, how are you?", # Replace with your actual message + # "name": "optional_name" # Only needed for named users in some APIs + } + ], + "max_tokens": 100, # Replace with your desired max tokens + "temperature": 0.7, # Replace with your desired temperature (0-2) + "top_p": 1, + "n": 1, + "stream": False, + "stop": None, # or ["string"] if you want to stop at certain phrases + "presence_penalty": 0, + "frequency_penalty": 0 + # Note: 'usage' is typically part of the response, not the request +} + +embeddingRequest = { + "model": "text-embedding-3-small", + "input": [ + "Message a embeder" + ], + "user": "string" +} + + +def test_read_main(): + response = client.get("/") + assert response.status_code == 404 + assert response.json() == {"detail":"Not Found"} +def test_list_models(): + response = client.get("/v1/models", headers={"Authorization": "Bearer " + token}) + assert response.status_code == 200 + data = response.json() + assert isinstance(data["data"], list) +def test_completion_chat(): + response = client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer " + token}, + json=chatRequest + ) + data = response.json() + # check type using pydantic class + assert ChatCompletionResponse(**data) + +def test_embedding_array(): + response = client.post( + "/v1/embeddings", + headers={"Authorization": "Bearer " + token}, + json=embeddingRequest + ,) + data = response.json() + assert EmbeddingResponse(**data) + +def test_embedding_string(): + embeddingRequest["input"] = "Message to embed" + response = client.post( + "/v1/embeddings", + headers={"Authorization": "Bearer " + token}, + json=embeddingRequest + ,) + data = response.json() + assert EmbeddingResponse(**data) + +# def test_completion_chat_b64Image(): +# chatRequest["messages"][0]["content"] = [ +# { +# "type": "text", +# "text": "What is image ?" +# }, +# { +# "type": "image_url", +# "image_url": {"url": b64_image} +# } +# ] +# print(chatRequest) +# response = client.post( +# "/v1/chat/completions", +# headers={"Authorization": "Bearer " + token}, +# json=chatRequest +# ) +# data = response.json() +# print(data) +# # check type using pydantic class +# assert ChatCompletionResponse(**data) + +# def test_completion_chat_urlImage(): +# chatRequest["messages"][0]["content"] = [ +# { +# "type": "text", +# "text": "What is image ?" +# }, +# { +# "type": "image_url", +# "image_url": {"url": url_image} +# } +# ] +# print(chatRequest) +# response = client.post( +# "/v1/chat/completions", +# headers={"Authorization": "Bearer " + token}, +# json=chatRequest +# ) +# data = response.json() +# print(data) +# # check type using pydantic class +# assert ChatCompletionResponse(**data) diff --git a/backend/tests.py b/backend/tests.py deleted file mode 100644 index 5aec45b..0000000 --- a/backend/tests.py +++ /dev/null @@ -1,80 +0,0 @@ -import requests -import os -import time -import base64 -from io import BytesIO - -class TestAPI: - api_url = os.getenv("API_URL", "http://localhost:8000") - api_token = os.getenv("API_TOKEN", "") - headers = {"Authorization": f"Bearer {api_token}"} if api_token else {} - system_prompt = "You are a helpful assistant." - user_message = "Hello, how can you assist me today?" - model = "gemma3" - image_url = "https://erasme.org/IMG/png/dataviz1.png" - - def health_check(self): - try: - response = requests.get(f"{self.api_url}/docs") - return response.status_code == 200 - except requests.ConnectionError: - return False - - def test_chat_completion(self): - payload = { - "model": self.model, - "messages": [ - {"role": "system", "content": self.system_prompt}, - {"role": "user", "content": self.user_message} - ] - } - response = requests.post(f"{self.api_url}/v1/chat/completions", json=payload, headers=self.headers) - data = response.json() - assert response.status_code == 200, f"Expected status code 200, got {response.status_code}" - - assert "choices" in data, "Response JSON does not contain 'choices'" - assert len(data["choices"]) > 0, "No choices returned in response" - print("Chat completion test passed.") - - def test_image_upload_b64(self): - # download image and convert to base64 - image_data = requests.get(self.image_url).content - image_b64 = base64.b64encode(image_data).decode('utf-8') - - payload = { - "model": self.model, - "messages": [ - {"role": "system", "content": self.system_prompt}, - { - "role": "user", - "content": [ - { - "type": "text", - "text": self.user_message - }, - { - "type": "image_url", - "image_url": { - "url": f"data:image/jpeg;base64,{image_b64}" - } - } - ] - } - ] - } - - response = requests.post(f"{self.api_url}/v1/chat/completions", json=payload, headers=self.headers) - data = response.json() - - assert response.status_code == 200, f"Expected status code 200, got {response.status_code}" - assert "choices" in data, "Response JSON does not contain 'choices'" - assert len(data["choices"]) > 0, "No choices returned in response" - print("Image upload (base64) test passed. with response:", data["choices"][0]["message"]["content"]) -if __name__ == "__main__": - api_tester = TestAPI() - while api_tester.health_check() is False: - print("Waiting for the API to be ready...") - time.sleep(1) - print("API is ready") - api_tester.test_chat_completion() - api_tester.test_image_upload_b64() \ No newline at end of file diff --git a/docker-compose.yaml b/docker-compose.yaml index d37f4a5..1e3c949 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -8,20 +8,17 @@ services: context: ./backend dockerfile: Dockerfile container_name: ai-proxy-test - depends_on: - - api volumes: - ./backend:/app + - ./config.yaml:/config.yaml + env_file: .env environment: - PYTHONUNBUFFERED=1 - - API_URL=http://api:8000 - - API_TOKEN=autobot3000 - command: ["python", "tests.py"] + command: ["pytest"] api: profiles: - api - - test image: erasme/ai-proxy:latest build: context: ./backend @@ -76,4 +73,4 @@ services: volumes: - grafana-storage: \ No newline at end of file + grafana-storage: