diff --git a/pyproject.toml b/pyproject.toml index 154db53..e172605 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,9 +95,13 @@ ml = [ "sentence-transformers>=2.2.0", # For VectorSearcher and DocstringIndexer "chromadb>=1.0.0", # Vector database for semantic search (1.0+ required for Cloud) ] +litellm = [ + "litellm>=1.80.0,<1.87.0", # LiteLLM AI gateway for 100+ LLM providers +] all = [ "sentence-transformers>=2.2.0", "chromadb>=1.0.0", + "litellm>=1.80.0,<1.87.0", ] [tool.ruff] diff --git a/src/kit/__init__.py b/src/kit/__init__.py index 36d9e54..1398dfd 100644 --- a/src/kit/__init__.py +++ b/src/kit/__init__.py @@ -16,7 +16,7 @@ from .vector_searcher import VectorSearcher try: - from .summaries import AnthropicConfig, GoogleConfig, LLMError, OpenAIConfig, Summarizer + from .summaries import AnthropicConfig, GoogleConfig, LiteLLMConfig, LLMError, OpenAIConfig, Summarizer except ImportError: # Allow kit to be imported even if LLM extras aren't installed. # Users will get an ImportError later if they try to use Summarizer. @@ -107,7 +107,7 @@ def _typer_make_metavar(self, ctx=None): # type: ignore[override] "get_tool_schemas", # Conditionally add Summarizer related classes if they were imported *( - ["Summarizer", "OpenAIConfig", "AnthropicConfig", "GoogleConfig", "LLMError"] + ["Summarizer", "OpenAIConfig", "AnthropicConfig", "GoogleConfig", "LiteLLMConfig", "LLMError"] if "Summarizer" in globals() else [] ), diff --git a/src/kit/cli.py b/src/kit/cli.py index 70045d3..f3bc79a 100644 --- a/src/kit/cli.py +++ b/src/kit/cli.py @@ -224,6 +224,8 @@ def commit_changes( ) elif detected_provider == LLMProvider.OLLAMA: new_api_key = "not_required" # Ollama doesn't need API key + elif detected_provider == LLMProvider.LITELLM: + new_api_key = os.getenv("LITELLM_API_KEY") or "litellm" else: # OpenAI new_api_key = os.getenv("KIT_OPENAI_TOKEN") or os.getenv("OPENAI_API_KEY") if not new_api_key: @@ -886,6 +888,8 @@ def review_pr( "Configuration error", "Set KIT_ANTHROPIC_TOKEN environment variable", ) + elif detected_provider == LLMProvider.LITELLM: + new_api_key = os.getenv("LITELLM_API_KEY") or "litellm" else: # OpenAI new_api_key = os.getenv("KIT_OPENAI_TOKEN") or os.getenv("OPENAI_API_KEY") if not new_api_key: @@ -1367,6 +1371,8 @@ def summarize_pr( ) elif detected_provider == LLMProvider.OLLAMA: new_api_key = "not_required" # Ollama doesn't need API key + elif detected_provider == LLMProvider.LITELLM: + new_api_key = os.getenv("LITELLM_API_KEY") or "litellm" else: # OpenAI new_api_key = os.getenv("KIT_OPENAI_TOKEN") or os.getenv("OPENAI_API_KEY") if not new_api_key: diff --git a/src/kit/deep_research.py b/src/kit/deep_research.py index b7e03b5..efda025 100644 --- a/src/kit/deep_research.py +++ b/src/kit/deep_research.py @@ -8,6 +8,7 @@ from kit.summaries import ( AnthropicConfig, GoogleConfig, + LiteLLMConfig, LLMError, OllamaConfig, OpenAIConfig, @@ -29,7 +30,7 @@ class ResearchResult: class DeepResearch: """LLM-based research for comprehensive answers.""" - def __init__(self, config: Optional[Union[OpenAIConfig, AnthropicConfig, GoogleConfig, OllamaConfig]] = None): + def __init__(self, config: Optional[Union[OpenAIConfig, AnthropicConfig, GoogleConfig, OllamaConfig, LiteLLMConfig]] = None): """Initialize with LLM config.""" self.config = config self._llm_client = None @@ -66,6 +67,14 @@ def _init_llm_client(self): elif isinstance(self.config, OllamaConfig): self._llm_client = "ollama" + elif isinstance(self.config, LiteLLMConfig): + try: + import litellm + + self._llm_client = litellm + except ImportError: + raise LLMError("litellm package not installed. Run: pip install 'cased-kit[litellm]'") + def research(self, query: str) -> ResearchResult: """ Perform research by prompting an LLM for a comprehensive answer. @@ -146,6 +155,26 @@ def research(self, query: str) -> ResearchResult: answer = response.json().get("response", "No response from Ollama") else: answer = f"Ollama error: {response.status_code}" + elif isinstance(self.config, LiteLLMConfig): + completion_kwargs = { + "model": self.config.model, + "messages": [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ], + "max_tokens": self.config.max_tokens, + "drop_params": self.config.drop_params, + } + if self.config.api_key: + completion_kwargs["api_key"] = self.config.api_key + if self.config.api_base: + completion_kwargs["api_base"] = self.config.api_base + + response = self._llm_client.completion(**completion_kwargs) + answer = response.choices[0].message.content + + if not answer: + answer = "The LLM returned an empty response. Please try again." else: answer = "No LLM configured." diff --git a/src/kit/llm_client_factory.py b/src/kit/llm_client_factory.py index 8c42206..1302765 100644 --- a/src/kit/llm_client_factory.py +++ b/src/kit/llm_client_factory.py @@ -10,7 +10,7 @@ if TYPE_CHECKING: from kit.pr_review.config import LLMConfig - from kit.summaries import AnthropicConfig, GoogleConfig, OllamaConfig, OpenAIConfig + from kit.summaries import AnthropicConfig, GoogleConfig, LiteLLMConfig, OllamaConfig, OpenAIConfig class LLMClientError(Exception): @@ -111,8 +111,32 @@ def create_ollama_client( return OllamaClient(base_url, model, session) +def create_litellm_client( + api_key: Optional[str] = None, + api_base: Optional[str] = None, +) -> Any: + """Create a LiteLLM client (returns the litellm module for direct use). + + Args: + api_key: Optional API key (passed per-call; provider env vars also work) + api_base: Optional base URL for LiteLLM proxy + + Returns: + The litellm module + + Raises: + LLMClientError: If the litellm package is not installed + """ + try: + import litellm + except ImportError: + raise LLMClientError("litellm package not installed. Run: pip install 'cased-kit[litellm]'") + + return litellm + + def create_client_from_config( - config: Union["OpenAIConfig", "AnthropicConfig", "GoogleConfig", "OllamaConfig"], + config: Union["OpenAIConfig", "AnthropicConfig", "GoogleConfig", "OllamaConfig", "LiteLLMConfig"], ) -> Any: """Create an LLM client from a summaries config object. @@ -127,7 +151,7 @@ def create_client_from_config( TypeError: If the config type is not recognized """ # Import here to avoid circular imports - from kit.summaries import AnthropicConfig, GoogleConfig, OllamaConfig, OpenAIConfig + from kit.summaries import AnthropicConfig, GoogleConfig, LiteLLMConfig, OllamaConfig, OpenAIConfig if isinstance(config, OpenAIConfig): return create_openai_client(config.api_key or "", config.base_url) @@ -137,6 +161,8 @@ def create_client_from_config( return create_google_client(config.api_key or "") elif isinstance(config, OllamaConfig): return create_ollama_client(config.base_url, config.model) + elif isinstance(config, LiteLLMConfig): + return create_litellm_client(config.api_key, config.api_base) else: raise TypeError(f"Unsupported config type: {type(config)}") @@ -173,5 +199,7 @@ def create_client_from_review_config( llm_config.model, session, ) + elif llm_config.provider == LLMProvider.LITELLM: + return create_litellm_client(llm_config.api_key, llm_config.api_base_url) else: raise ValueError(f"Unsupported LLM provider: {llm_config.provider}") diff --git a/src/kit/pr_review/agentic_reviewer.py b/src/kit/pr_review/agentic_reviewer.py index 5ebcce7..dd40bf2 100644 --- a/src/kit/pr_review/agentic_reviewer.py +++ b/src/kit/pr_review/agentic_reviewer.py @@ -924,6 +924,121 @@ async def make_api_call(): return "Analysis completed after maximum turns" + async def _run_agentic_analysis_litellm(self, initial_prompt: str) -> str: + """Run multi-turn agentic analysis using LiteLLM (OpenAI-compatible tool calling).""" + try: + import litellm + except ImportError: + raise RuntimeError("litellm package not installed. Run: pip install 'cased-kit[litellm]'") + + tools = [ + { + "type": "function", + "function": { + "name": tool["name"], + "description": tool["description"], + "parameters": tool["input_schema"], + }, + } + for tool in self._get_available_tools() + ] + messages: List[Dict[str, Any]] = [{"role": "user", "content": initial_prompt}] + + max_turns = self.max_turns + turn = 0 + + while turn < max_turns: + turn += 1 + print(f"\U0001f916 Agentic turn {turn}...") + + if turn >= max_turns - 3: + messages.append( + { + "role": "user", + "content": f"URGENT: You are on turn {turn} of {max_turns}. You MUST finalize your review NOW using the finalize_review tool. Do not use any other tools.", + } + ) + elif turn >= self.finalize_threshold: + messages.append( + { + "role": "user", + "content": f"You are on turn {turn} of {max_turns}. Please finalize your review soon using the finalize_review tool with your comprehensive analysis.", + } + ) + + try: + completion_kwargs: Dict[str, Any] = { + "model": self.config.llm.model, + "tools": tools, + "messages": messages, + "max_tokens": self.config.llm.max_tokens, + "drop_params": True, + } + if self.config.llm.api_key and self.config.llm.api_key != "litellm": + completion_kwargs["api_key"] = self.config.llm.api_key + if self.config.llm.api_base_url: + completion_kwargs["api_base"] = self.config.llm.api_base_url + + async def make_api_call(): + import asyncio + + loop = asyncio.get_event_loop() + return await loop.run_in_executor( + None, + lambda: litellm.completion(**completion_kwargs), + ) + + response = await retry_with_backoff(make_api_call) + + # Track cost + input_tokens = getattr(response.usage, "prompt_tokens", 0) if response.usage else 0 + output_tokens = getattr(response.usage, "completion_tokens", 0) if response.usage else 0 + self.cost_tracker.track_llm_usage( + self.config.llm.provider, self.config.llm.model, input_tokens, output_tokens + ) + + message = response.choices[0].message + messages.append({"role": "assistant", "content": message.content, "tool_calls": message.tool_calls}) + + if message.tool_calls: + print( + f"\U0001f680 Executing {len(message.tool_calls)} {'tool' if len(message.tool_calls) == 1 else 'tools'} in parallel..." + ) + + tool_tasks = [] + tool_call_info = [] + + for tool_call in message.tool_calls: + tool_name = tool_call.function.name + tool_input = json.loads(tool_call.function.arguments) + + if tool_name == "finalize_review": + final_review = tool_input.get("review_markdown", "") + if final_review: + return final_review + return tool_input.get("analysis", "No analysis provided") + + tool_tasks.append(self._execute_tool(tool_name, tool_input)) + tool_call_info.append((tool_call.id, tool_name)) + + results = await asyncio.gather(*tool_tasks, return_exceptions=True) + + for (tool_id, tool_name), result in zip(tool_call_info, results): + if isinstance(result, Exception): + result_str = f"Error: {result}" + else: + result_str = str(result)[:4000] + messages.append({"role": "tool", "tool_call_id": tool_id, "content": result_str}) + else: + text_content = message.content or "" + if text_content and turn > 1: + return text_content + + except Exception as e: + return f"Error during agentic analysis turn {turn}: {e}" + + return "Analysis completed after maximum turns" + async def analyze_pr_agentic(self, repo_path: str, pr_details: Dict[str, Any], files: List[Dict[str, Any]]) -> str: """Run agentic analysis of the PR.""" from kit import Repository @@ -1031,6 +1146,8 @@ async def analyze_pr_agentic(self, repo_path: str, pr_details: Dict[str, Any], f "Please use --provider anthropic, openai, or google for agentic reviews, " "or run without --agentic for standard reviews with Ollama." ) + elif self.config.llm.provider == LLMProvider.LITELLM: + analysis = await self._run_agentic_analysis_litellm(initial_prompt) else: analysis = await self._run_agentic_analysis_openai(initial_prompt) diff --git a/src/kit/pr_review/commit_generator.py b/src/kit/pr_review/commit_generator.py index 542a8a3..a6c7955 100644 --- a/src/kit/pr_review/commit_generator.py +++ b/src/kit/pr_review/commit_generator.py @@ -130,6 +130,8 @@ async def analyze_changes_for_commit(self, repo_path: str) -> str: message = await self._generate_with_google(commit_prompt) elif self.config.llm.provider == LLMProvider.OLLAMA: message = await self._generate_with_ollama(commit_prompt) + elif self.config.llm.provider == LLMProvider.LITELLM: + message = await self._generate_with_litellm(commit_prompt) else: message = await self._generate_with_openai(commit_prompt) @@ -252,6 +254,40 @@ async def _generate_with_ollama(self, prompt: str) -> str: except Exception as e: return f"Update files (error: {e})" + async def _generate_with_litellm(self, prompt: str) -> str: + """Generate commit message using LiteLLM.""" + try: + import litellm + except ImportError: + return "Update files (error: litellm not installed)" + + try: + completion_kwargs: Dict[str, Any] = { + "model": self.config.llm.model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": 200, + "drop_params": True, + } + if self.config.llm.api_key and self.config.llm.api_key != "litellm": + completion_kwargs["api_key"] = self.config.llm.api_key + if self.config.llm.api_base_url: + completion_kwargs["api_base"] = self.config.llm.api_base_url + + response = await asyncio.to_thread(litellm.completion, **completion_kwargs) + + # Track cost + input_tokens = getattr(response.usage, "prompt_tokens", 0) if response.usage else 0 + output_tokens = getattr(response.usage, "completion_tokens", 0) if response.usage else 0 + self.cost_tracker.track_llm_usage( + self.config.llm.provider, self.config.llm.model, input_tokens, output_tokens + ) + + content = response.choices[0].message.content + return content if content is not None else "Update files" + + except Exception as e: + return f"Update files (error: {e})" + def commit_with_message(self, message: str) -> None: """Execute git commit with the generated message.""" try: diff --git a/src/kit/pr_review/config.py b/src/kit/pr_review/config.py index 85ae826..7165019 100644 --- a/src/kit/pr_review/config.py +++ b/src/kit/pr_review/config.py @@ -16,6 +16,7 @@ class LLMProvider(Enum): OPENAI = "openai" OLLAMA = "ollama" GOOGLE = "google" + LITELLM = "litellm" class ReviewDepth(Enum): @@ -69,6 +70,10 @@ def _detect_provider_from_model(model_name: str) -> Optional[LLMProvider]: if any(pattern in stripped_model for pattern in google_patterns): return LLMProvider.GOOGLE + # LiteLLM model patterns - models with litellm routing prefix + if model_lower.startswith("litellm/") or model_lower.startswith("litellm_proxy/"): + return LLMProvider.LITELLM + # Ollama model patterns - popular models available in Ollama ollama_patterns = [ "llama", @@ -238,7 +243,15 @@ def from_file( raise ValueError(f"Unsupported LLM provider: {provider_str}. Use: {[p.value for p in LLMProvider]}") # Default models and API key environment variables - if provider == LLMProvider.ANTHROPIC: + if provider == LLMProvider.LITELLM: + default_model = "openai/gpt-4" + api_key_env = "LITELLM_API_KEY (or provider-specific env vars like OPENAI_API_KEY)" + config_api_key = llm_data.get("api_key") + if _is_placeholder_token(config_api_key): + config_api_key = None + # LiteLLM can read provider env vars directly, so api_key is optional + api_key = config_api_key or os.getenv("LITELLM_API_KEY") or "litellm" + elif provider == LLMProvider.ANTHROPIC: default_model = "claude-sonnet-4-5" api_key_env = "KIT_ANTHROPIC_TOKEN or ANTHROPIC_API_KEY" config_api_key = llm_data.get("api_key") @@ -265,8 +278,8 @@ def from_file( config_api_key = None # Treat placeholder as missing api_key = config_api_key or os.getenv("KIT_OPENAI_TOKEN") or os.getenv("OPENAI_API_KEY") - # Ollama doesn't require API keys, so skip validation for it - if not api_key and provider != LLMProvider.OLLAMA: + # Ollama and LiteLLM don't require API keys, so skip validation for them + if not api_key and provider not in (LLMProvider.OLLAMA, LLMProvider.LITELLM): # Provide a more helpful error message based on the provider if provider == LLMProvider.ANTHROPIC and not (config_data.get("llm") or {}).get("provider"): # User didn't specify a provider, defaulted to Anthropic @@ -410,6 +423,14 @@ def create_default_config_file(self, config_path: Optional[str] = None) -> str: # Add a commented-out example for local Ollama usage local_ollama_example = """\ +# Example LiteLLM configuration (access 100+ LLM providers via single interface): +# llm: +# provider: litellm +# model: "anthropic/claude-sonnet-4-5" # Or "openai/gpt-4", "bedrock/anthropic.claude-3-sonnet", etc. +# api_key: "" # Optional - LiteLLM reads provider env vars (OPENAI_API_KEY, ANTHROPIC_API_KEY, etc.) +# api_base_url: "http://localhost:4000" # Optional - only needed for LiteLLM proxy server +# max_tokens: 4000 + # Example Ollama configuration (completely free local AI): # llm: # provider: ollama diff --git a/src/kit/pr_review/reviewer.py b/src/kit/pr_review/reviewer.py index baa358b..6118c5e 100644 --- a/src/kit/pr_review/reviewer.py +++ b/src/kit/pr_review/reviewer.py @@ -183,6 +183,8 @@ async def analyze_pr_with_kit(self, repo_path: str, pr_details: Dict[str, Any], analysis = await self._analyze_with_google_enhanced(analysis_prompt) elif self.config.llm.provider == LLMProvider.OLLAMA: analysis = await self._analyze_with_ollama_enhanced(analysis_prompt) + elif self.config.llm.provider == LLMProvider.LITELLM: + analysis = await self._analyze_with_litellm_enhanced(analysis_prompt) else: analysis = await self._analyze_with_openai_enhanced(analysis_prompt) @@ -336,6 +338,40 @@ async def _analyze_with_ollama_enhanced(self, enhanced_prompt: str) -> str: except Exception as e: return f"Error during enhanced Ollama analysis: {e}" + async def _analyze_with_litellm_enhanced(self, enhanced_prompt: str) -> str: + """Analyze using LiteLLM with enhanced kit context.""" + try: + import litellm + except ImportError: + return "Error: litellm package not installed. Run: pip install 'cased-kit[litellm]'" + + try: + completion_kwargs: Dict[str, Any] = { + "model": self.config.llm.model, + "messages": [{"role": "user", "content": enhanced_prompt}], + "max_tokens": self.config.llm.max_tokens, + "drop_params": True, + } + if self.config.llm.api_key and self.config.llm.api_key != "litellm": + completion_kwargs["api_key"] = self.config.llm.api_key + if self.config.llm.api_base_url: + completion_kwargs["api_base"] = self.config.llm.api_base_url + + response = await asyncio.to_thread(litellm.completion, **completion_kwargs) + + # Track cost + input_tokens = getattr(response.usage, "prompt_tokens", 0) if response.usage else 0 + output_tokens = getattr(response.usage, "completion_tokens", 0) if response.usage else 0 + self.cost_tracker.track_llm_usage( + self.config.llm.provider, self.config.llm.model, input_tokens, output_tokens + ) + + content = response.choices[0].message.content + return content if content is not None else "No response content" + + except Exception as e: + return f"Error during enhanced LLM analysis: {e}" + def review_pr(self, pr_input: str) -> str: """Review a PR with intelligent analysis.""" try: @@ -788,6 +824,8 @@ async def analyze_local_diff_with_kit( analysis = await self._analyze_with_google_enhanced(analysis_prompt) elif self.config.llm.provider == LLMProvider.OLLAMA: analysis = await self._analyze_with_ollama_enhanced(analysis_prompt) + elif self.config.llm.provider == LLMProvider.LITELLM: + analysis = await self._analyze_with_litellm_enhanced(analysis_prompt) else: analysis = await self._analyze_with_openai_enhanced(analysis_prompt) diff --git a/src/kit/summaries.py b/src/kit/summaries.py index 5945e3e..71350a5 100644 --- a/src/kit/summaries.py +++ b/src/kit/summaries.py @@ -115,6 +115,27 @@ def __post_init__(self): self.base_url = self.base_url.rstrip("/") +@dataclass +class LiteLLMConfig: + """Configuration for LiteLLM AI gateway access. + + LiteLLM provides a unified interface to 100+ LLM providers (OpenAI, Anthropic, + Google, AWS Bedrock, Azure, Cohere, etc.) through a single API. + + Model format uses LiteLLM's routing syntax, e.g.: + - "openai/gpt-4" or just "gpt-4" + - "anthropic/claude-3-opus-20240229" + - "bedrock/anthropic.claude-3-sonnet-20240229-v1:0" + - "vertex_ai/gemini-pro" + """ + + model: str = "openai/gpt-4" + api_key: Optional[str] = field(default=None) + api_base: Optional[str] = field(default=None) + max_tokens: int = 1000 + drop_params: bool = True + + # todo: make configurable MAX_CODE_LENGTH_CHARS = 50000 # Max characters for a single function/class summary MAX_FILE_SUMMARIZE_CHARS = 25000 # Max characters for file content in summarize_file @@ -157,7 +178,7 @@ class Summarizer: """Provides methods to summarize code using a configured LLM.""" _tokenizer_cache: Dict[str, Any] = {} # Cache for tiktoken encoders - config: Optional[Union[OpenAIConfig, AnthropicConfig, GoogleConfig, OllamaConfig]] + config: Optional[Union[OpenAIConfig, AnthropicConfig, GoogleConfig, OllamaConfig, LiteLLMConfig]] repo: "Repository" _llm_client: Optional[Any] # type: ignore @@ -305,7 +326,7 @@ def _count_openai_chat_tokens(self, messages: List[Dict[str, str]], model_name: def __init__( self, repo: "Repository", - config: Optional[Union[OpenAIConfig, AnthropicConfig, GoogleConfig, OllamaConfig]] = None, + config: Optional[Union[OpenAIConfig, AnthropicConfig, GoogleConfig, OllamaConfig, LiteLLMConfig]] = None, llm_client: Optional[Any] = None, ): """ @@ -510,6 +531,31 @@ def summarize_file(self, file_path: str) -> str: except Exception as e: logger.warning(f"Ollama API error for {file_path}: {e}") summary = f"Summary generation failed: Ollama API error ({e})" + elif isinstance(self.config, LiteLLMConfig): + try: + import litellm + except ImportError: + raise LLMError("litellm package not installed. Run: pip install 'cased-kit[litellm]'") + + messages_for_api = [ + {"role": "system", "content": system_prompt_text}, + {"role": "user", "content": user_prompt_text}, + ] + completion_kwargs: Dict[str, Any] = { + "model": self.config.model, + "messages": messages_for_api, + "max_tokens": self.config.max_tokens, + "drop_params": self.config.drop_params, + } + if self.config.api_key: + completion_kwargs["api_key"] = self.config.api_key + if self.config.api_base: + completion_kwargs["api_base"] = self.config.api_base + + response = litellm.completion(**completion_kwargs) + summary = response.choices[0].message.content + if response.usage: + logger.debug(f"LiteLLM API usage for {file_path}: {response.usage}") else: # This should never happen with our current logic, but as a safeguard raise LLMError(f"Unsupported LLM configuration type: {type(self.config) if self.config else None}") @@ -679,6 +725,31 @@ def summarize_function(self, file_path: str, function_name: str) -> str: except Exception as e: logger.warning(f"Ollama API error for {function_name} in {file_path}: {e}") summary = f"Summary generation failed: Ollama API error ({e})" + elif isinstance(self.config, LiteLLMConfig): + try: + import litellm + except ImportError: + raise LLMError("litellm package not installed. Run: pip install 'cased-kit[litellm]'") + + messages_for_api = [ + {"role": "system", "content": system_prompt_text}, + {"role": "user", "content": user_prompt_text}, + ] + completion_kwargs: Dict[str, Any] = { + "model": self.config.model, + "messages": messages_for_api, + "max_tokens": self.config.max_tokens, + "drop_params": self.config.drop_params, + } + if self.config.api_key: + completion_kwargs["api_key"] = self.config.api_key + if self.config.api_base: + completion_kwargs["api_base"] = self.config.api_base + + response = litellm.completion(**completion_kwargs) + summary = response.choices[0].message.content + if response.usage: + logger.debug(f"LiteLLM API usage for {function_name} in {file_path}: {response.usage}") else: # This should never happen with our current logic, but as a safeguard raise LLMError(f"Unsupported LLM configuration type: {type(self.config) if self.config else None}") @@ -846,6 +917,31 @@ def summarize_class(self, file_path: str, class_name: str) -> str: except Exception as e: logger.warning(f"Ollama API error for {class_name} in {file_path}: {e}") summary = f"Summary generation failed: Ollama API error ({e})" + elif isinstance(self.config, LiteLLMConfig): + try: + import litellm + except ImportError: + raise LLMError("litellm package not installed. Run: pip install 'cased-kit[litellm]'") + + messages_for_api = [ + {"role": "system", "content": system_prompt_text}, + {"role": "user", "content": user_prompt_text}, + ] + completion_kwargs: Dict[str, Any] = { + "model": self.config.model, + "messages": messages_for_api, + "max_tokens": self.config.max_tokens, + "drop_params": self.config.drop_params, + } + if self.config.api_key: + completion_kwargs["api_key"] = self.config.api_key + if self.config.api_base: + completion_kwargs["api_base"] = self.config.api_base + + response = litellm.completion(**completion_kwargs) + summary = response.choices[0].message.content + if response.usage: + logger.debug(f"LiteLLM API usage for {class_name} in {file_path}: {response.usage}") else: # This should never happen with our current logic, but as a safeguard raise LLMError(f"Unsupported LLM configuration type: {type(self.config) if self.config else None}") diff --git a/tests/test_llm_client_factory.py b/tests/test_llm_client_factory.py index 80146d1..58bacfb 100644 --- a/tests/test_llm_client_factory.py +++ b/tests/test_llm_client_factory.py @@ -9,6 +9,7 @@ create_client_from_config, create_client_from_review_config, create_google_client, + create_litellm_client, create_ollama_client, create_openai_client, ) @@ -236,3 +237,74 @@ def test_creates_ollama_client_from_review_config(self): mock_create.assert_called_once_with("http://localhost:11434", "llama3", mock_session) assert result == mock_client + + +class TestCreateLiteLLMClient: + """Tests for create_litellm_client function.""" + + def test_creates_client(self): + """Test creating LiteLLM client returns the litellm module.""" + with patch.dict("sys.modules", {"litellm": MagicMock()}): + result = create_litellm_client() + assert result is not None + + def test_creates_client_with_params(self): + """Test creating LiteLLM client with api_key and api_base.""" + with patch.dict("sys.modules", {"litellm": MagicMock()}): + result = create_litellm_client(api_key="test-key", api_base="http://localhost:4000") + assert result is not None + + +class TestCreateLiteLLMClientFromConfig: + """Tests for create_client_from_config with LiteLLMConfig.""" + + def test_creates_litellm_client_from_config(self): + """Test creating client from LiteLLMConfig.""" + from kit.summaries import LiteLLMConfig + + with patch("kit.llm_client_factory.create_litellm_client") as mock_create: + mock_client = MagicMock() + mock_create.return_value = mock_client + + config = LiteLLMConfig(model="openai/gpt-4", api_key="test-key", api_base="http://localhost:4000") + result = create_client_from_config(config) + + mock_create.assert_called_once_with("test-key", "http://localhost:4000") + assert result == mock_client + + def test_creates_litellm_client_from_config_no_key(self): + """Test creating client from LiteLLMConfig without api_key.""" + from kit.summaries import LiteLLMConfig + + with patch("kit.llm_client_factory.create_litellm_client") as mock_create: + mock_client = MagicMock() + mock_create.return_value = mock_client + + config = LiteLLMConfig(model="anthropic/claude-3-opus-20240229") + result = create_client_from_config(config) + + mock_create.assert_called_once_with(None, None) + assert result == mock_client + + +class TestCreateLiteLLMClientFromReviewConfig: + """Tests for create_client_from_review_config with LITELLM provider.""" + + def test_creates_litellm_client_from_review_config(self): + """Test creating client from LLMConfig with LITELLM provider.""" + from kit.pr_review.config import LLMConfig, LLMProvider + + with patch("kit.llm_client_factory.create_litellm_client") as mock_create: + mock_client = MagicMock() + mock_create.return_value = mock_client + + config = LLMConfig( + provider=LLMProvider.LITELLM, + api_key="test-key", + model="openai/gpt-4", + api_base_url="http://localhost:4000", + ) + result = create_client_from_review_config(config) + + mock_create.assert_called_once_with("test-key", "http://localhost:4000") + assert result == mock_client