diff --git a/__init__.py b/__init__.py index 611c9dd..124376e 100644 --- a/__init__.py +++ b/__init__.py @@ -3,6 +3,6 @@ llm_client: A Python package for interacting with LLM models through Ollama. """ from _llm._llm.llm import LLM -from _llm._llm.tool_registy import register_tool, get_tools +from _llm._llm.tool_registry import register_tool, get_tools __all__ = ["LLM", "register_tool", "get_tools"] \ No newline at end of file diff --git a/_llm/__init__.py b/_llm/__init__.py index 97a5d0b..64d80bb 100644 --- a/_llm/__init__.py +++ b/_llm/__init__.py @@ -1,6 +1,6 @@ from .llm import LLM # re-export the class from the module -from .tool_registy import register_tool, get_tools +from .tool_registry import register_tool, get_tools # Define public API __all__ = ["LLM", "register_tool", "get_tools"] \ No newline at end of file diff --git a/_llm/llm.py b/_llm/llm.py index b2bd712..aa029e5 100644 --- a/_llm/llm.py +++ b/_llm/llm.py @@ -9,19 +9,19 @@ from openai.types.chat import ChatCompletion, ChatCompletionChunk from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_message import ChatCompletionMessage from openai.types.responses import ParsedResponse - import backoff import env_manager import json try: - from .tool_registy import get_tools, register_tool + from .tool_registry import get_tools, parse_function_call_arguments, execute_tool except ImportError: - from tool_registy import get_tools, register_tool + + from _llm._llm.tool_registry import get_tools, parse_function_call_arguments, execute_tool try: from colorprinter.print_color import * except ImportError: - from colorprinter.print_color import * + from colorprinter.colorprinter.print_color import * @@ -487,6 +487,7 @@ class LLM: # Call the OpenAI API else: response: ChatCompletion = self.client.chat.completions.create(**kwargs) + # Try to extract backend information if available try: response_headers = getattr(response, "_headers", {}) @@ -720,10 +721,64 @@ class LLM: else: choice = response.choices[0] message: ChatCompletionMessage = choice.message + print(message) + + + if hasattr(message, 'tool_calls') and message.tool_calls: + # Hantera flera verktygsanrop sequensielt + for tool_call in message.tool_calls: + try: + fn = getattr(tool_call, "function", None) or (tool_call.get("function") if isinstance(tool_call, dict) else None) + if not fn: + continue + func_name = getattr(fn, "name", None) or (fn.get("name") if isinstance(fn, dict) else None) + raw_args = getattr(fn, "arguments", None) or (fn.get("arguments") if isinstance(fn, dict) else None) + + # Automatisk JSON-parsing av argument om de kommer som sträng (för vLLM-kompatibilitet) + if isinstance(raw_args, str): + try: + parsed_args = json.loads(raw_args) + # Uppdatera function.arguments med den parsade versionen för enklare användning + if hasattr(fn, "arguments"): + fn.arguments = parsed_args # Uppdatera objektet direkt + elif isinstance(fn, dict): + fn["arguments"] = parsed_args + except json.JSONDecodeError as e: + print_red(f"Warning: Could not parse tool arguments as JSON: {e}") + # Fallback till parse_function_call_arguments för robusthet + parsed_args = parse_function_call_arguments(raw_args) + else: + parsed_args = raw_args if isinstance(raw_args, dict) else {} + + # Kör verktyget via tool_registry.execute_tool (validering och typ-coercion görs där) + tool_result = execute_tool(func_name, parsed_args) + # Sätt in tool-result i messages så modellen kan läsa det vidare + tool_content = tool_result if isinstance(tool_result, str) else json.dumps(tool_result, ensure_ascii=False) + self.messages.append({"role": "tool", "name": func_name, "content": tool_content}) + except Exception as e: + print_red(f"Error executing tool {func_name}: {e}") + # append error to messages so model sees it (and you can debug) + self.messages.append({"role": "tool", "name": func_name or "unknown", "content": json.dumps({"error": str(e)})}) + # fallback: older SDKs / shapes: + if hasattr(message, 'function_call') and message.function_call: + fc = message.function_call + func_name = getattr(fc, "name", None) or (fc.get("name") if isinstance(fc, dict) else None) + args_raw = getattr(fc, "arguments", None) or (fc.get("arguments") if isinstance(fc, dict) else None) + parsed_args = parse_function_call_arguments(args_raw) + try: + tool_result = execute_tool(func_name, parsed_args) + tool_content = tool_result if isinstance(tool_result, str) else json.dumps(tool_result, ensure_ascii=False) + self.messages.append({"role": "tool", "name": func_name, "content": tool_content}) + except Exception as e: + self.messages.append({"role": "tool", "name": func_name or "unknown", "content": json.dumps({"error": str(e)})}) + + + # Hämta textsvar från meddelandet result: str = message.content if hasattr(message, 'content_text'): result: str = message.content_text - # Store in message history (without tool calls for clean history) + + # Spara i meddelandehistorik (utan verktygsanrop för ren historik) self.messages.append({"role": "assistant", "content": result}) if not self.chat: self.messages = [self.messages[0]] @@ -1120,7 +1175,6 @@ if __name__ == "__main__": print(response.__dict__) response = llm.generate("What's the weather like in San Francisco? Also calculate 15 * 7 for me.", model='vllm') print(response.__dict__) - exit() # Define a tool for calculations @register_tool @@ -1198,7 +1252,7 @@ if __name__ == "__main__": except Exception as e: print(f"❌ Tools test failed: {e}") - # Test 3: Thinking mode (use vllm model since reasoning model doesn't exist) + # Test 3: Thinking mode (use vLLM model since reasoning model doesn't exist) print("\n3️⃣ Thinking Mode Test (using vllm)") print("-" * 30) try: diff --git a/_llm/ollama-cloud-test.py b/_llm/ollama-cloud-test.py new file mode 100644 index 0000000..fa9afb3 --- /dev/null +++ b/_llm/ollama-cloud-test.py @@ -0,0 +1,13 @@ +from ollama import Client + +client = Client() + +messages = [ + { + 'role': 'user', + 'content': 'Why is the sky blue?', + }, +] + +for part in client.chat('gpt-oss:120b-cloud', messages=messages, stream=True): + print(part['message']['content'], end='', flush=True) \ No newline at end of file diff --git a/_llm/tool_registry.py b/_llm/tool_registry.py new file mode 100644 index 0000000..a7f52d3 --- /dev/null +++ b/_llm/tool_registry.py @@ -0,0 +1,176 @@ +import inspect, json, re, ast +from typing import Callable, Dict, Any, List, get_origin, get_args +from pydantic import BaseModel + +TOOL_REGISTRY: Dict[str, Dict[str, Any]] = {} + +# --- type mapping --- +def _pytype_to_jsonschema(t): + origin = get_origin(t) + if origin is list or origin is List: + args = get_args(t) + item_type = args[0] if args else str + return {"type": "array", "items": _pytype_to_jsonschema(item_type)} + if inspect.isclass(t) and issubclass(t, BaseModel): + sch = t.schema() + return {"type": "object", **sch} + mapping = { + str: {"type": "string"}, + int: {"type": "integer"}, + float: {"type": "number"}, + bool: {"type": "boolean"}, + dict: {"type": "object"}, + list: {"type": "array", "items": {"type": "string"}}, + } + return mapping.get(t, {"type": "string"}) + +# --- docstring parser (Google style) --- +def _parse_google_docstring(docstring: str): + if not docstring: + return {"description": "", "params": {}} + lines = [ln.rstrip() for ln in docstring.splitlines()] + desc_lines = [] + i = 0 + while i < len(lines) and not lines[i].lower().startswith(("args:", "arguments:")): + if lines[i].strip(): + desc_lines.append(lines[i].strip()) + i += 1 + description = " ".join(desc_lines).strip() + params = {} + if i < len(lines): + i += 1 + while i < len(lines): + line = lines[i].strip() + if not line: + i += 1 + continue + m = re.match(r'^(\w+)\s*(?:\(([^)]+)\))?\s*:\s*(.*)$', line) + if m: + name = m.group(1) + desc = m.group(3) + j = i + 1 + while j < len(lines) and not re.match(r'^\w+\s*(?:\([^)]+\))?\s*:', lines[j].strip()): + if lines[j].strip(): + desc += " " + lines[j].strip() + j += 1 + params[name] = {"description": desc.strip(), "type": m.group(2)} + i = j + continue + i += 1 + return {"description": description, "params": params} + +# --- helper: make OpenAI-style function spec --- +def _wrap_openai_function_schema(name: str, description: str, parameters: dict): + """Create OpenAI function calling format with 'function' wrapper""" + params = parameters.copy() + if params.get("type") != "object": + params = {"type": "object", "properties": params.get("properties", params), "required": params.get("required", [])} + params.setdefault("additionalProperties", False) + + # Return in OpenAI function calling format with 'function' wrapper + return { + "type": "function", + "function": { + "name": name, + "description": description, + "parameters": params + } + } + +# --- decorator to register tools --- +def register_tool(func: Callable = None, *, name: str = None, description: str = None, schema: dict = None): + def _register(f): + fname = name or f.__name__ + doc = _parse_google_docstring(f.__doc__) + func_description = description or doc["description"] or "" + if schema is not None: + func_schema = schema + else: + sig = inspect.signature(f) + props = {} + required = [] + for param_name, param in sig.parameters.items(): + ann = param.annotation if param.annotation is not inspect._empty else str + prop_schema = _pytype_to_jsonschema(ann) + if param_name in doc["params"]: + prop_schema["description"] = doc["params"][param_name]["description"] + props[param_name] = prop_schema + if param.default is inspect._empty: + required.append(param_name) + func_schema = {"type": "object", "properties": props, "required": required, "additionalProperties": False} + TOOL_REGISTRY[fname] = { + "callable": f, + "schema": _wrap_openai_function_schema(fname, func_description, func_schema) + } + return f + if func is None: + return _register + else: + return _register(func) + +# --- what to send to model --- +def get_tools() -> List[dict]: + """Return OpenAI-compatible functions list with proper 'function' wrapper.""" + return [entry["schema"] for entry in TOOL_REGISTRY.values()] + +# --- robust parser for arguments --- +def parse_function_call_arguments(raw) -> dict: + if isinstance(raw, dict): + return raw + if not isinstance(raw, str): + return {"_raw_unexpected": str(type(raw)), "value": raw} + try: + return json.loads(raw) + except json.JSONDecodeError: + pass + try: + return ast.literal_eval(raw) + except Exception: + pass + stripped = raw.strip() + if re.match(r'^(SELECT|WITH)\b', stripped, flags=re.IGNORECASE): + return {"sql_query": stripped} + m = re.search(r'\{.*\}', raw, flags=re.DOTALL) + if m: + candidate = m.group(0) + try: + return json.loads(candidate) + except Exception: + try: + return ast.literal_eval(candidate) + except Exception: + pass + return {"_raw": raw} + +# --- safe executor --- +def execute_tool(name: str, args: dict): + """ + Execute registered callable with args (basic validation). + Returns Python object (dict/list/str). + """ + entry = TOOL_REGISTRY.get(name) + if not entry: + raise RuntimeError(f"Function {name} not registered") + fn = entry["callable"] + # simple SQL safety example: if function expects sql_query ensure SELECT + if "sql_query" in args: + q = args["sql_query"].strip() + if not re.match(r'^(SELECT|WITH)\b', q, flags=re.IGNORECASE): + raise ValueError("Only SELECT/ WITH queries allowed in sql_query") + if q.endswith(";"): + args["sql_query"] = q[:-1] + # Prepare kwargs with minimal type coercion + sig = inspect.signature(fn) + kwargs = {} + for pname, param in sig.parameters.items(): + if pname not in args: + continue + val = args[pname] + ann = param.annotation if param.annotation is not inspect._empty else None + origin = get_origin(ann) + if origin in (list, List) and isinstance(val, str): + kwargs[pname] = [x.strip() for x in val.split(",") if x.strip() != ""] + else: + kwargs[pname] = val + result = fn(**kwargs) + return result diff --git a/_llm/tool_registy.py b/_llm/tool_registy.py deleted file mode 100644 index 3eadaa9..0000000 --- a/_llm/tool_registy.py +++ /dev/null @@ -1,214 +0,0 @@ -# assume your client already has: import inspect, json -from typing import Callable, Dict, Any -import inspect, json -import re -from pydantic import BaseModel - -TOOL_REGISTRY: Dict[str, Dict[str, Any]] = {} - -def _parse_google_docstring(docstring: str) -> Dict[str, Any]: - """Parse Google-style docstring to extract description and parameter info.""" - if not docstring: - return {"description": "", "params": {}} - - # Split into lines and clean up - lines = [line.strip() for line in docstring.strip().split('\n')] - - # Find the main description (everything before Args:) - description_lines = [] - i = 0 - while i < len(lines): - if lines[i].lower().startswith('args:') or lines[i].lower().startswith('arguments:'): - break - description_lines.append(lines[i]) - i += 1 - - description = ' '.join(description_lines).strip() - - # Parse parameters section - params = {} - if i < len(lines): - i += 1 # Skip the "Args:" line - while i < len(lines): - line = lines[i] - if line.lower().startswith(('returns:', 'yields:', 'raises:', 'note:', 'example:')): - break - - # Match parameter format: param_name (type): description - match = re.match(r'^\s*(\w+)\s*(?:\(([^)]+)\))?\s*:\s*(.*)$', line) - if match: - param_name = match.group(1) - param_type = match.group(2) - param_desc = match.group(3) - - # Collect multi-line descriptions - j = i + 1 - while j < len(lines) and lines[j] and not re.match(r'^\s*\w+\s*(?:\([^)]+\))?\s*:', lines[j]): - param_desc += ' ' + lines[j].strip() - j += 1 - - params[param_name] = { - "description": param_desc.strip(), - "type": param_type.strip() if param_type else None - } - i = j - 1 - - i += 1 - - return {"description": description, "params": params} - -def _pytype_to_jsonschema(t): - # Very-small helper; extend as needed or use pydantic models for complex types - mapping = {str: {"type": "string"}, int: {"type": "integer"}, - float: {"type": "number"}, bool: {"type": "boolean"}, - dict: {"type": "object"}, list: {"type": "array"}} - return mapping.get(t, {"type": "string"}) # fallback to string - -def register_tool(func: Callable = None, *, name: str = None, description: str = None, schema: dict = None): - """ - Use as decorator or call directly: - @register_tool - def foo(x: int): ... - or - register_tool(func=myfunc, name="myfunc", schema=...) - """ - def _register(f): - fname = name or f.__name__ - - # Parse docstring for description and parameter info - docstring_info = _parse_google_docstring(f.__doc__) - func_description = description or docstring_info["description"] or "" - - # If explicit schema provided, use it - if schema is not None: - func_schema = schema - else: - sig = inspect.signature(f) - props = {} - required = [] - for param_name, param in sig.parameters.items(): - ann = param.annotation - # If user used a Pydantic BaseModel as a single arg, use its schema - if inspect.isclass(ann) and issubclass(ann, BaseModel): - func_schema = ann.schema() - # wrap into a single-arg object if necessary - props = func_schema.get("properties", {}) - required = func_schema.get("required", []) - # done early - for single-model param - break - - # Create property schema from type annotation - prop_schema = _pytype_to_jsonschema(ann) - - # Add description from docstring if available - if param_name in docstring_info["params"]: - prop_schema["description"] = docstring_info["params"][param_name]["description"] - - props[param_name] = prop_schema - if param.default is inspect._empty: - required.append(param_name) - - if 'func_schema' not in locals(): - func_schema = { - "type": "object", - "properties": props, - "required": required - } - - TOOL_REGISTRY[fname] = { - "callable": f, - "schema": { - "type": "function", - "function": { - "name": fname, - "description": func_description, - "parameters": func_schema - } - } - } - return f - - if func is None: - return _register - else: - return _register(func) - -def get_tools() -> list: - """Return list of function schemas (JSON) to send to the model""" - return [v["schema"] for v in TOOL_REGISTRY.values()] - -def handle_function_call_and_inject_result(response_choice, messages): - """ - Given the model choice (response.choices[0]) and your messages list: - - extracts function/tool call - - executes the registered python callable - - appends the tool result as a tool message and returns it - """ - # Support different shapes: some SDKs use .message.tool_calls, others .message.function_call - msg = getattr(response_choice, "message", None) or (response_choice.get("message") if isinstance(response_choice, dict) else None) - func_name = None - func_args = None - # try tool_calls style - if msg: - tool_calls = getattr(msg, "tool_calls", None) or (msg.get("tool_calls") if isinstance(msg, dict) else None) - if tool_calls: - tc = tool_calls[0] - fn = getattr(tc, "function", None) or (tc.get("function") if isinstance(tc, dict) else None) - func_name = getattr(fn, "name", None) or (fn.get("name") if isinstance(fn, dict) else None) - func_args = getattr(fn, "arguments", None) or (fn.get("arguments") if isinstance(fn, dict) else None) - # fallback to function_call - if func_name is None: - fc = getattr(msg, "function_call", None) or (msg.get("function_call") if isinstance(msg, dict) else None) - if fc: - func_name = getattr(fc, "name", None) or fc.get("name") - args_raw = getattr(fc, "arguments", None) or fc.get("arguments") - # arguments are often a JSON string depending on SDK shape - if isinstance(args_raw, str): - try: - func_args = json.loads(args_raw) - except Exception: - func_args = None - else: - func_args = args_raw - - if not func_name: - return None # no function call found - - entry = TOOL_REGISTRY.get(func_name) - if not entry: - raise RuntimeError(f"Function {func_name} not registered") - - result = entry["callable"](**(func_args or {})) - # convert result to string/JSON for tool message - tool_content = result if isinstance(result, str) else json.dumps(result) - # append tool message so model can see the result - messages.append({"role": "tool", "name": func_name, "content": tool_content}) - return tool_content - -if __name__ == "__main__": - # Example usage and test - @register_tool - def add(x: int, y: int) -> int: - """Add two integers - Args: - x (int): First integer - y (int): Second integer - Returns: - int: Sum of x and y - """ - return x + y - - @register_tool(name="echo", description="Echoes the input string") - def echo_message(message: str) -> str: - """Echo the input message - Args: - message (str): The message to echo - Returns: - str: The echoed message - """ - return message - - print("Registered tools:") - import pprint - for info in get_tools(): - pprint.pprint(info) \ No newline at end of file