llm_studio_client.py

import aiohttp
import asyncio
import logging
from typing import List, Dict, Any

logger = logging.getLogger(__name__)

class LMStudioMessage:
    def __init__(self, content: str):
        self.content = content

class LMStudioClient:
    def __init__(self, model: str, base_url: str = "http://localhost:1234/v1", temperature: float = 0.7, max_tokens: int = -1):
        self.model = model
        self.base_url = base_url
        self.temperature = temperature
        self.max_tokens = max_tokens
    
    async def ainvoke(self, messages: List[Dict[str, str]]) -> LMStudioMessage:
        """Async invoke method compatible with LangChain interface"""
        try:
            async with aiohttp.ClientSession() as session:
                payload = {
                    "model": self.model,
                    "messages": messages,
                    "temperature": self.temperature,
                    "max_tokens": self.max_tokens,
                    "stream": False
                }
                
                logger.debug(f"Sending request to {self.base_url}/chat/completions with model: {self.model}")
                
                async with session.post(
                    f"{self.base_url}/chat/completions",
                    json=payload,
                    headers={"Content-Type": "application/json"},
                    timeout=aiohttp.ClientTimeout(total=30)
                ) as response:
                    if response.status == 200:
                        data = await response.json()
                        content = data["choices"][0]["message"]["content"]
                        return LMStudioMessage(content)
                    else:
                        error_text = await response.text()
                        logger.error(f"LM Studio API error {response.status}: {error_text}")
                        raise Exception(f"LM Studio API error: {response.status} - {error_text}")
        except aiohttp.ClientError as e:
            logger.error(f"Connection error to LM Studio: {e}")
            raise Exception(f"LM Studio connection error: {e}")
        except asyncio.TimeoutError:
            logger.error("Request to LM Studio timed out")
            raise Exception("LM Studio request timeout")

    def invoke(self, messages: List[Dict[str, str]]) -> LMStudioMessage:
        """Sync invoke method for compatibility"""
        return asyncio.run(self.ainvoke(messages))