echodict/llama.cpp

1import argparse2import json3import requests4import logging5import sys6 7handler = logging.StreamHandler(sys.stdout)8handler.terminator = ""   # ← no newline9logging.basicConfig(level=logging.INFO, format='%(message)s', handlers=[handler])10logger = logging.getLogger("server-test-model")11 12 13def run_query(url, messages, tools=None, stream=False, tool_choice=None):14    payload = {15        "messages": messages,16        "stream": stream,17        "max_tokens": 5000,18    }19    if tools:20        payload["tools"] = tools21    if tool_choice:22        payload["tool_choice"] = tool_choice23 24    try:25        response = requests.post(url, json=payload, stream=stream)26        response.raise_for_status()27    except requests.exceptions.RequestException as e:28        if e.response is not None:29            logger.info(f"Response error: {e} for {e.response.content}\n")30        else:31            logger.info(f"Error connecting to server: {e}\n")32        return None33 34    full_content = ""35    reasoning_content = ""36    tool_calls = []37 38    if stream:39        logger.info(f"--- Streaming response (Tools: {bool(tools)}) ---\n")40        for line in response.iter_lines():41            if line:42                decoded_line = line.decode("utf-8")43                if decoded_line.startswith("data: "):44                    data_str = decoded_line[6:]45                    if data_str == "[DONE]":46                        break47                    try:48                        data = json.loads(data_str)49                        if "choices" in data and len(data["choices"]) > 0:50                            delta = data["choices"][0].get("delta", {})51 52                            # Content53                            content_chunk = delta.get("content", "")54                            if content_chunk:55                                full_content += content_chunk56                                logger.info(content_chunk)57 58                            # Reasoning59                            reasoning_chunk = delta.get("reasoning_content", "")60                            if reasoning_chunk:61                                reasoning_content += reasoning_chunk62                                logger.info(f"\x1B[3m{reasoning_chunk}\x1B[0m")63 64                            # Tool calls65                            if "tool_calls" in delta:66                                for tc in delta["tool_calls"]:67                                    index = tc.get("index")68                                    if index is not None:69                                        while len(tool_calls) <= index:70                                            # Using "function" as type default but could be flexible71                                            tool_calls.append(72                                                {73                                                    "id": "",74                                                    "type": "function",75                                                    "function": {76                                                        "name": "",77                                                        "arguments": "",78                                                    },79                                                }80                                            )81 82                                        if "id" in tc:83                                            tool_calls[index]["id"] += tc["id"]84                                        if "function" in tc:85                                            if "name" in tc["function"]:86                                                tool_calls[index]["function"][87                                                    "name"88                                                ] += tc["function"]["name"]89                                            if "arguments" in tc["function"]:90                                                tool_calls[index]["function"][91                                                    "arguments"92                                                ] += tc["function"]["arguments"]93 94                    except json.JSONDecodeError:95                        logger.info(f"Failed to decode JSON: {data_str}\n")96        logger.info("\n--- End of Stream ---\n")97    else:98        logger.info(f"--- Non-streaming response (Tools: {bool(tools)}) ---\n")99        data = response.json()100        if "choices" in data and len(data["choices"]) > 0:101            message = data["choices"][0].get("message", {})102            full_content = message.get("content", "")103            reasoning_content = message.get("reasoning_content", "")104            tool_calls = message.get("tool_calls", [])105            logger.info(full_content)106        logger.info("--- End of Response ---\n")107 108    return {109        "content": full_content,110        "reasoning_content": reasoning_content,111        "tool_calls": tool_calls,112    }113 114 115def test_chat(url, stream):116    logger.info(f"\n=== Testing Chat (Stream={stream}) ===\n")117    messages = [{"role": "user", "content": "What is the capital of France?"}]118    result = run_query(url, messages, stream=stream)119 120    if result:121        if result["content"]:122            logger.info("PASS: Output received.\n")123        else:124            logger.info("WARN: No content received (valid if strict tool call, but unexpected here).\n")125 126        if result.get("reasoning_content"):127            logger.info(f"INFO: Reasoning content detected ({len(result['reasoning_content'])} chars).\n")128        else:129            logger.info("INFO: No reasoning content detected (Standard model behavior).\n")130    else:131        logger.info("FAIL: No result.\n")132 133 134def test_tool_call(url, stream):135    logger.info(f"\n=== Testing Tool Call (Stream={stream}) ===\n")136    messages = [137        {138            "role": "user",139            "content": "What is the weather in London? Please use the get_weather tool.",140        }141    ]142    tools = [143        {144            "type": "function",145            "function": {146                "name": "get_weather",147                "description": "Get the current weather in a given location",148                "parameters": {149                    "type": "object",150                    "properties": {151                        "location": {152                            "type": "string",153                            "description": "The city and state, e.g. San Francisco, CA",154                        },155                        "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},156                    },157                    "required": ["location"],158                },159            },160        }161    ]162 163    result = run_query(url, messages, tools=tools, tool_choice="auto", stream=stream)164 165    if result:166        tcs = result.get("tool_calls")167        if tcs and len(tcs) > 0:168            logger.info("PASS: Tool calls detected.")169            for tc in tcs:170                func = tc.get("function", {})171                logger.info(f"  Tool: {func.get('name')}, Args: {func.get('arguments')}\n")172        else:173            logger.info(f"FAIL: No tool calls. Content: {result['content']}\n")174 175        if result.get("reasoning_content"):176            logger.info(177                f"INFO: Reasoning content detected during tool call ({len(result['reasoning_content'])} chars).\n"178            )179    else:180        logger.info("FAIL: Query failed.\n")181 182 183def main():184    parser = argparse.ArgumentParser(description="Test llama-server functionality.")185    parser.add_argument("--host", default="localhost", help="Server host")186    parser.add_argument("--port", default=8080, type=int, help="Server port")187    args = parser.parse_args()188 189    base_url = f"http://{args.host}:{args.port}/v1/chat/completions"190    logger.info(f"Testing server at {base_url}\n")191 192    # Non-streaming tests193    test_chat(base_url, stream=False)194    test_tool_call(base_url, stream=False)195 196    # Streaming tests197    test_chat(base_url, stream=True)198    test_tool_call(base_url, stream=True)199 200 201if __name__ == "__main__":202    main()203