src/api/client.py

#
# SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org>
# SPDX-License-Identifier: Apache-2.0
#

import aiohttp
import json
from config import (
    ENDPOINT,
    API_KEY,
    MODEL,
    STREAM,
    AIOHTTP,
    RETRY
)
from ..tools.mapping import TOOLS

async def client(messages):
    async with aiohttp.ClientSession(
        connector=aiohttp.TCPConnector(**AIOHTTP),
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}"
        }
    ) as session:
        for attempt in range(RETRY):
            async with session.post(
                ENDPOINT,
                json={
                    "model": MODEL,
                    "messages": messages,
                    "tools": TOOLS,
                    "tool_choice": "auto",
                    "stream": STREAM
                }
            ) as response:
                if response.status != 200:
                    if attempt == RETRY - 1:
                        error_message = await response.text()
                        raise Exception(f"Error ({response.status}): {error_message}")
                    continue
                
                buffer = ""
                
                async for parts in response.content.iter_any():
                    if not parts:
                        continue
                    
                    buffer += parts.decode('utf-8')
                    
                    while '\n' in buffer:
                        line, buffer = buffer.split('\n', 1)
                        data = line.strip()
                        
                        if not data:
                            continue
                        
                        if data.startswith("data: "):
                            data = data[6:]
                        
                        if data == "[DONE]":
                            return
                        
                        if data:
                            try:
                                chunk = json.loads(data)
                                yield chunk
                            except json.JSONDecodeError:
                                continue
                
                return