"""LLM client using vLLM for Llama-3.1-8B-Instruct inference with tool calling."""
import json, os, time, traceback
from typing import Dict, List, Optional, Any
from openai import OpenAI


class LLMClient:
    """OpenAI-compatible client for vLLM-served Llama model."""
    def __init__(self, model="meta-llama/Llama-3.1-8B-Instruct", base_url="http://localhost:8000/v1", api_key="dummy", temperature=0.001, max_tokens=1024):
        self.model, self.temperature, self.max_tokens = model, temperature, max_tokens
        self.client = OpenAI(base_url=base_url, api_key=api_key)
        self.conversation_history = []

    def change_messages(self, messages):
        self.conversation_history = messages

    def parse(self, tools, process_id=0, max_retries=3):
        use_messages = []
        for msg in self.conversation_history:
            if msg.get("valid", True) != False:
                clean_msg = {k: v for k, v in msg.items() if k != "valid"}
                clean_msg.pop("function_call", None)
                use_messages.append(clean_msg)
        for attempt in range(max_retries):
            try:
                kwargs = {"model": self.model, "messages": use_messages, "max_tokens": self.max_tokens, "temperature": self.temperature, "frequency_penalty": 0, "presence_penalty": 0}
                if tools and any(t["function"]["name"] != "Finish" for t in tools):
                    kwargs["tools"] = tools
                    kwargs["parallel_tool_calls"] = False
                response = self.client.chat.completions.create(**kwargs)
                message = response.choices[0].message
                total_tokens = response.usage.total_tokens if response.usage else 0
                msg_dict = {"role": "assistant", "content": message.content}
                if message.tool_calls:
                    msg_dict["tool_calls"] = [{"id": tc.id, "type": "function", "function": {"name": tc.function.name, "arguments": tc.function.arguments}} for tc in message.tool_calls]
                if process_id == 0: print(f"[process({process_id})] tokens: {total_tokens}")
                return msg_dict, 0, total_tokens
            except Exception as e:
                print(f"[process({process_id})] Attempt {attempt+1} error: {repr(e)}")
                traceback.print_exc()
                if attempt < max_retries - 1: time.sleep(2)
        return {"role": "assistant", "content": "Error generating response"}, -1, 0