switchforge 1.0.0__tar.gz → 1.1.0__tar.gz

{switchforge-1.0.0 → switchforge-1.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: switchforge
-Version: 1.0.0
+Version: 1.1.0
 Summary: AI-powered backend test generation engine
 Project-URL: Homepage, https://theswitchcompany.online/products/forge/core
 Project-URL: Repository, https://github.com/switchcompany/forge-core
@@ -18,18 +18,19 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Software Development :: Testing
 Requires-Python: >=3.10
 Requires-Dist: httpx>=0.27.0
-Requires-Dist: instructor>=1.3.0
-Requires-Dist: litellm>=1.40.0
+Requires-Dist: openai>=1.30.0
 Requires-Dist: pydantic>=2.7.0
 Requires-Dist: pyyaml>=6.0
 Requires-Dist: rich>=13.7.0
-Requires-Dist: tiktoken>=0.7.0
 Requires-Dist: typer>=0.12.0
+Provides-Extra: accurate-tokens
+Requires-Dist: tiktoken>=0.7.0; extra == 'accurate-tokens'
 Provides-Extra: dev
 Requires-Dist: pyinstaller>=6.0; extra == 'dev'
 Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
 Requires-Dist: pytest>=8.0; extra == 'dev'
 Requires-Dist: ruff>=0.5.0; extra == 'dev'
+Requires-Dist: tiktoken>=0.7.0; extra == 'dev'
 Description-Content-Type: text/markdown
 
 # Switchforge (Forge Core CLI)

{switchforge-1.0.0 → switchforge-1.1.0}/forge_core/ai/provider.py

@@ -1,27 +1,44 @@
-"""LiteLLM-based AI provider — unified interface for 100+ models."""
+"""AI provider — direct OpenAI SDK, zero native dependencies."""
 
 from __future__ import annotations
 
+import os
 from typing import Any
 
-import litellm
+from openai import OpenAI
 
 from forge_core.models.config import AIConfig, AIProvider
 from forge_core.utils import logger
 from forge_core.utils.tokens import count_tokens
 
-# Suppress LiteLLM's verbose logging
-litellm.suppress_debug_info = True
+
+def _get_client(config: AIConfig) -> OpenAI:
+    """Create an OpenAI-compatible client for any provider."""
+    api_key = config.api_key or os.environ.get("OPENAI_API_KEY", "")
+
+    if config.provider == AIProvider.ANTHROPIC:
+        base_url = config.base_url or "https://api.anthropic.com/v1"
+        api_key = config.api_key or os.environ.get("ANTHROPIC_API_KEY", "")
+    elif config.provider == AIProvider.OLLAMA:
+        base_url = config.base_url or "http://localhost:11434/v1"
+        api_key = "ollama"
+    elif config.provider == AIProvider.AZURE:
+        base_url = config.base_url or os.environ.get("AZURE_OPENAI_ENDPOINT", "")
+        api_key = config.api_key or os.environ.get("AZURE_OPENAI_API_KEY", "")
+    elif config.base_url:
+        base_url = config.base_url
+    else:
+        base_url = None  # default OpenAI
+
+    kwargs: dict[str, Any] = {"api_key": api_key}
+    if base_url:
+        kwargs["base_url"] = base_url
+
+    return OpenAI(**kwargs)
 
 
 def _resolve_model(config: AIConfig) -> str:
-    """Resolve the model string for LiteLLM based on provider."""
-    if config.provider == AIProvider.OLLAMA:
-        return f"ollama/{config.model}"
-    if config.provider == AIProvider.AZURE:
-        return f"azure/{config.model}"
-    if config.provider == AIProvider.BEDROCK:
-        return f"bedrock/{config.model}"
+    """Resolve the model name."""
     return config.model
 
 
@@ -32,19 +49,10 @@ def complete(
     json_mode: bool = False,
     max_tokens: int | None = None,
 ) -> str:
-    """Send a completion request to the configured AI provider.
-
-    Args:
-        config: AI provider configuration.
-        system_prompt: System-level instructions.
-        user_prompt: User message with code/context.
-        json_mode: If True, request JSON response format.
-        max_tokens: Override max tokens for this call.
-
-    Returns:
-        The AI's response text.
-    """
+    """Send a completion request to the configured AI provider."""
     model = _resolve_model(config)
+    client = _get_client(config)
+
     messages = [
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": user_prompt},
@@ -57,19 +65,14 @@ def complete(
         "max_tokens": max_tokens or config.max_tokens,
     }
 
-    if config.api_key:
-        kwargs["api_key"] = config.api_key
-    if config.base_url:
-        kwargs["api_base"] = config.base_url
     if json_mode:
         kwargs["response_format"] = {"type": "json_object"}
 
-    # Log token usage
     input_tokens = count_tokens(system_prompt + user_prompt, config.model)
     logger.info(f"AI call → {model} ({input_tokens} input tokens)")
 
     try:
-        response = litellm.completion(**kwargs)
+        response = client.chat.completions.create(**kwargs)
         content = response.choices[0].message.content or ""
 
         output_tokens = count_tokens(content, config.model)
@@ -88,11 +91,9 @@ def complete_with_fallback(
     fallback_models: list[str] | None = None,
     json_mode: bool = False,
 ) -> str:
-    """Try primary model, fall back to alternatives on failure.
-
-    Useful for complex classes where one model might fail.
-    """
+    """Try primary model, fall back to alternatives on failure."""
     models = [_resolve_model(config)] + (fallback_models or [])
+    client = _get_client(config)
 
     last_error = None
     for model in models:
@@ -106,12 +107,10 @@ def complete_with_fallback(
                 "temperature": config.temperature,
                 "max_tokens": config.max_tokens,
             }
-            if config.api_key:
-                kwargs["api_key"] = config.api_key
             if json_mode:
                 kwargs["response_format"] = {"type": "json_object"}
 
-            response = litellm.completion(**kwargs)
+            response = client.chat.completions.create(**kwargs)
             return response.choices[0].message.content or ""
         except Exception as e:
             last_error = e

switchforge-1.1.0/forge_core/ai/structured.py

@@ -0,0 +1,62 @@
+"""Structured AI outputs using OpenAI JSON mode + pydantic parsing."""
+
+from __future__ import annotations
+
+import json
+from typing import TypeVar
+
+from pydantic import BaseModel
+
+from forge_core.models.config import AIConfig
+from forge_core.ai.provider import complete
+from forge_core.utils import logger
+
+T = TypeVar("T", bound=BaseModel)
+
+
+def extract(
+    config: AIConfig,
+    system_prompt: str,
+    user_prompt: str,
+    response_model: type[T],
+    max_retries: int = 2,
+) -> T:
+    """Extract structured data from AI using JSON mode + pydantic.
+
+    Forces the AI to return data matching a pydantic model schema.
+    Retries on validation failure.
+    """
+    schema = response_model.model_json_schema()
+    structured_prompt = (
+        f"{system_prompt}\n\n"
+        f"You MUST respond with valid JSON matching this schema:\n"
+        f"```json\n{json.dumps(schema, indent=2)}\n```\n"
+        f"Respond ONLY with the JSON object, no other text."
+    )
+
+    logger.info(f"Structured extraction → {config.model} → {response_model.__name__}")
+
+    last_error = None
+    for attempt in range(max_retries + 1):
+        try:
+            raw = complete(config, structured_prompt, user_prompt, json_mode=True)
+            # Strip markdown code fences if present
+            cleaned = raw.strip()
+            if cleaned.startswith("```"):
+                cleaned = cleaned.split("\n", 1)[1] if "\n" in cleaned else cleaned
+                if cleaned.endswith("```"):
+                    cleaned = cleaned[:-3]
+                cleaned = cleaned.strip()
+
+            result = response_model.model_validate_json(cleaned)
+            logger.success(f"Extracted {response_model.__name__} successfully")
+            return result
+        except Exception as e:
+            last_error = e
+            if attempt < max_retries:
+                logger.warn(f"Extraction attempt {attempt + 1} failed: {e}, retrying...")
+            continue
+
+    raise RuntimeError(
+        f"Structured extraction failed after {max_retries + 1} attempts: {last_error}"
+    )

{switchforge-1.0.0 → switchforge-1.1.0}/forge_core/utils/tokens.py

@@ -1,12 +1,23 @@
-"""Token counting and prompt size management using tiktoken."""
+"""Token counting and prompt size management."""
 
 from __future__ import annotations
 
-import tiktoken
+try:
+    import tiktoken
+    _HAS_TIKTOKEN = True
+except ImportError:
+    _HAS_TIKTOKEN = False
+
+
+def _estimate_tokens(text: str) -> int:
+    """Rough token estimate (~4 chars per token) when tiktoken is unavailable."""
+    return len(text) // 4 + 1
 
 
 def count_tokens(text: str, model: str = "gpt-4o") -> int:
     """Count tokens in text for a given model."""
+    if not _HAS_TIKTOKEN:
+        return _estimate_tokens(text)
     try:
         enc = tiktoken.encoding_for_model(model)
     except KeyError:
@@ -16,6 +27,9 @@ def count_tokens(text: str, model: str = "gpt-4o") -> int:
 
 def truncate_to_tokens(text: str, max_tokens: int, model: str = "gpt-4o") -> str:
     """Truncate text to fit within max_tokens."""
+    if not _HAS_TIKTOKEN:
+        char_limit = max_tokens * 4
+        return text[:char_limit] if len(text) > char_limit else text
     try:
         enc = tiktoken.encoding_for_model(model)
     except KeyError:

{switchforge-1.0.0 → switchforge-1.1.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "switchforge"
-version = "1.0.0"
+version = "1.1.0"
 description = "AI-powered backend test generation engine"
 readme = "README.md"
 license = "MIT"
@@ -25,22 +25,22 @@ classifiers = [
 ]
 
 dependencies = [
-    "litellm>=1.40.0",
-    "instructor>=1.3.0",
+    "openai>=1.30.0",
     "typer>=0.12.0",
     "rich>=13.7.0",
     "pyyaml>=6.0",
-    "tiktoken>=0.7.0",
     "httpx>=0.27.0",
     "pydantic>=2.7.0",
 ]
 
 [project.optional-dependencies]
+accurate-tokens = ["tiktoken>=0.7.0"]
 dev = [
     "pytest>=8.0",
     "pytest-asyncio>=0.23",
     "ruff>=0.5.0",
     "pyinstaller>=6.0",
+    "tiktoken>=0.7.0",
 ]
 
 [project.scripts]

switchforge-1.0.0/forge_core/ai/structured.py

@@ -1,68 +0,0 @@
-"""Structured AI outputs using instructor + pydantic."""
-
-from __future__ import annotations
-
-from typing import Any, TypeVar
-
-import instructor
-import litellm
-from pydantic import BaseModel
-
-from forge_core.models.config import AIConfig
-from forge_core.ai.provider import _resolve_model
-from forge_core.utils import logger
-
-T = TypeVar("T", bound=BaseModel)
-
-
-def extract(
-    config: AIConfig,
-    system_prompt: str,
-    user_prompt: str,
-    response_model: type[T],
-    max_retries: int = 2,
-) -> T:
-    """Extract structured data from AI using instructor.
-
-    Forces the AI to return data matching a pydantic model schema.
-    Retries on validation failure.
-
-    Args:
-        config: AI provider configuration.
-        system_prompt: System-level instructions.
-        user_prompt: User message with code/context.
-        response_model: Pydantic model class to extract.
-        max_retries: Number of retries on validation failure.
-
-    Returns:
-        An instance of response_model populated by the AI.
-    """
-    model = _resolve_model(config)
-
-    client = instructor.from_litellm(litellm.completion)
-
-    kwargs: dict[str, Any] = {
-        "model": model,
-        "messages": [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_prompt},
-        ],
-        "temperature": config.temperature,
-        "max_retries": max_retries,
-        "response_model": response_model,
-    }
-
-    if config.api_key:
-        kwargs["api_key"] = config.api_key
-    if config.base_url:
-        kwargs["api_base"] = config.base_url
-
-    logger.info(f"Structured extraction → {model} → {response_model.__name__}")
-
-    try:
-        result = client.chat.completions.create(**kwargs)
-        logger.success(f"Extracted {response_model.__name__} successfully")
-        return result
-    except Exception as e:
-        logger.error(f"Structured extraction failed: {e}")
-        raise