Spaces:
Running
Running
| from langchain_anthropic import ChatAnthropic | |
| from langchain_openai import ChatOpenAI | |
| from langchain_ollama import ChatOllama | |
| from langchain_google_genai import ChatGoogleGenerativeAI | |
| from langchain_core.language_models.base import BaseLanguageModel | |
| from langchain_core.messages import BaseMessage, HumanMessage, AIMessage | |
| from typing import Optional, Dict, List, Any, Sequence | |
| from langchain.tools import BaseTool | |
| import os | |
| import requests | |
| import json | |
| from dotenv import load_dotenv | |
| from dataclasses import dataclass | |
| import ollama | |
| import copy | |
| load_dotenv() | |
| class LLMProvider: | |
| def __init__(self, api_keys: Optional[Dict[str, str]] = None): | |
| self.api_keys = {} | |
| if api_keys: | |
| self.api_keys = api_keys | |
| self.providers: Dict[str, Any] = {} | |
| self._setup_providers() | |
| def _get_ollama_models(self) -> List[str]: | |
| """Get list of available Ollama models using the ollama package""" | |
| try: | |
| models = ollama.list() | |
| return [model.model for model in models['models']] | |
| except Exception: | |
| return [] | |
| def _setup_providers(self): | |
| # Google Gemini | |
| if google_key := os.getenv('GOOGLE_API_KEY') or self.api_keys.get('Gemimi'): | |
| self.providers['Gemini'] = ChatGoogleGenerativeAI( | |
| model="gemini-1.5-pro-latest", | |
| api_key=google_key) | |
| # Anthropicsel | |
| if anthropic_key := os.getenv('ANTHROPIC_API_KEY') or self.api_keys.get('Claude'): | |
| self.providers['Claude'] = ChatAnthropic( | |
| api_key=anthropic_key, | |
| model_name="claude-3-5-sonnet-20241022", | |
| ) | |
| # OpenAI | |
| if openai_key := os.getenv('OPENAI_API_KEY') or self.api_keys.get('ChatGPT'): | |
| self.providers['ChatGPT'] = ChatOpenAI( | |
| api_key=openai_key, | |
| model_name="gpt-4o-2024-11-20", | |
| max_completion_tokens=4096, | |
| ) | |
| # Ollama (local) | |
| try: | |
| # Get available Ollama models using the ollama package | |
| ollama_models = self._get_ollama_models() | |
| for model in ollama_models: | |
| self.providers[f'Ollama-{model}'] = ChatOllama(model=model) | |
| except Exception: | |
| pass # Ollama not available | |
| def get_available_providers(self) -> list[str]: | |
| """Return list of available provider names""" | |
| return list(self.providers.keys()) | |
| def get_provider(self, name: str) -> Optional[Any]: | |
| """Get LLM provider by name""" | |
| return self.providers.get(name) | |