main.py

import asyncio
import json
from litellm import acompletion

from functions import web_search, execute_commands

def load_json_file(path: str):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

prompts = load_json_file("./prompts.json")
tools_schema = prompts.get("tool_schema")

def translate_api_type(api_type: str):
    if api_type == "openai-completions":
        return "openai"
    else:
        return api_type
async def run_async_agent(user_request: str, model_config: dict, max_steps: int = 5):
    if not model_config:
        return

    messages = [
        {"role": "system", "content": prompts.get("agents")},
        {"role": "user", "content": user_request},
    ]

    litellm_model_name = f"{translate_api_type(model_config['api_type'])}/{model_config['model_id']}"
    brave_api_key = model_config.get("web_search", {}).get("api_key")

    for step in range(max_steps):
        print(f"\n--- Step {step + 1}/{max_steps} ---")

        response = await acompletion(
            model=litellm_model_name,
            api_base=model_config["base_url"],
            api_key=model_config["api_key"],
            max_tokens=model_config["max_tokens"],
            messages=messages,
            tools=tools_schema,
            tool_choice="auto",
        )

        response_message = response.choices[0].message
        messages.append(response_message.model_dump(exclude_none=True))

        if response_message.tool_calls:
            for tool_call in response_message.tool_calls:
                args = json.loads(tool_call.function.arguments)

                if tool_call.function.name == "execute_commands":
                    results = await execute_commands(args.get("type"), args.get("command"), prompts)
                elif tool_call.function.name == "web_search":
                    results = await web_search(args.get("query"), brave_api_key)
                else:
                    results = [{"error": f"Unknown tool: {tool_call.function.name}"}]

                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "name": tool_call.function.name,
                        "content": json.dumps(results, ensure_ascii=False),
                    }
                )
            continue

        print("\n[Final Answer]", response_message.content)
        return


def load_app_config(config_path: str = "config.json"):
    try:
        return load_json_file(config_path)
    except FileNotFoundError:
        print(f"Config file not found: {config_path}")
        return None
    except Exception as e:
        print(f"Failed to load config: {e}")
        return None


def load_model_config(config_path: str = "config.json"):
    try:
        config_data = load_app_config(config_path)
        if not config_data:
            return None

        app_config = config_data.get("app", {})
        provider_name = app_config.get("provider")
        target_model_id = app_config.get("model")

        if not provider_name:
            raise ValueError("Missing config: app.provider")
        if not target_model_id:
            raise ValueError("Missing config: app.model")

        provider_config = config_data.get(provider_name)
        if not provider_config:
            raise ValueError(f"Provider not found: {provider_name}")

        for model in provider_config.get("models", []):
            if model.get("id") == target_model_id:
                return {
                    "base_url": provider_config.get("baseUrl"),
                    "api_type": provider_config.get("api"),
                    "model_id": model.get("id"),
                    "api_key": model.get("api-token"),
                    "max_tokens": model.get("maxTokens"),
                    "web_search": model.get("web", {}).get("search", {}),
                    "app": app_config,
                }

        raise ValueError(f"Model not found under {provider_name}: {target_model_id}")
    except Exception as e:
        print(f"Failed to load model config: {e}")
        return None


if __name__ == "__main__":
    my_config = load_model_config()
    print(my_config)

    if my_config:
        app_config = my_config.get("app", {})
        asyncio.run(
            run_async_agent(
                app_config.get("user_request", input("Enter your task:\n>")),
                my_config,
                max_steps=app_config.get("max_steps", 20),
            )
        )
    else:
        print("Unable to load model config, exiting.")