Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(weave): Add model to provider info map to source code, add make command for updating #2825

Merged
merged 3 commits into from
Oct 31, 2024
Merged

fix(weave): Add model to provider info map to source code, add make command for updating #2825

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e1d0c0d
use tmp as store path for model providers file
KyleGoyette Oct 30, 2024
849451f
add model providers file, modify makefile, update server code
KyleGoyette Oct 30, 2024
98abdd0
Merge branch 'master' into kyle/fix-model-providers-path
raubitsj Oct 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,5 +17,4 @@ gha-creds-*.json
.coverage
.nox
*.log
*/file::memory:?cache=shared
weave/trace_server/model_providers
*/file::memory:?cache=shared
3 changes: 3 additions & 0 deletions weave/Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,2 +1,5 @@
update_costs:
python trace_server/costs/update_costs.py

update_model_providers:
python trace_server/model_providers/model_providers.py
7 changes: 2 additions & 5 deletions weave/trace_server/clickhouse_trace_server_batched.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,8 +81,7 @@
from weave.trace_server.ids import generate_id
from weave.trace_server.llm_completion import lite_llm_completion
from weave.trace_server.model_providers.model_providers import (
MODEL_PROVIDERS_FILE,
fetch_model_to_provider_info_map,
read_model_to_provider_info_map,
)
from weave.trace_server.orm import ParamBuilder, Row
from weave.trace_server.secret_fetcher_context import _secret_fetcher_context
Expand Down Expand Up @@ -191,9 +190,7 @@ def __init__(
self._flush_immediately = True
self._call_batch: list[list[Any]] = []
self._use_async_insert = use_async_insert
self._model_to_provider_info_map = fetch_model_to_provider_info_map(
MODEL_PROVIDERS_FILE
)
self._model_to_provider_info_map = read_model_to_provider_info_map()

@classmethod
def from_env(cls, use_async_insert: bool = False) -> "ClickHouseTraceServer":
Expand Down
1 change: 1 addition & 0 deletions weave/trace_server/model_providers/model_providers.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
{"gpt-4": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4o": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4o-audio-preview": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4o-audio-preview-2024-10-01": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4o-mini": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4o-mini-2024-07-18": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "o1-mini": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "o1-mini-2024-09-12": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "o1-preview": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "o1-preview-2024-09-12": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "chatgpt-4o-latest": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4o-2024-05-13": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4o-2024-08-06": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-turbo-preview": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-0314": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-0613": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-32k": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-32k-0314": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-32k-0613": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-turbo": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-turbo-2024-04-09": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-1106-preview": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-0125-preview": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-vision-preview": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-4-1106-vision-preview": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-3.5-turbo": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-3.5-turbo-0301": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-3.5-turbo-0613": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-3.5-turbo-1106": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-3.5-turbo-0125": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-3.5-turbo-16k": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "gpt-3.5-turbo-16k-0613": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "ft:gpt-3.5-turbo": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "ft:gpt-3.5-turbo-0125": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "ft:gpt-3.5-turbo-1106": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "ft:gpt-3.5-turbo-0613": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "ft:gpt-4-0613": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "ft:gpt-4o-2024-08-06": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "ft:gpt-4o-mini-2024-07-18": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "text-embedding-3-large": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "text-embedding-3-small": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "text-embedding-ada-002": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "text-embedding-ada-002-v2": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "text-moderation-stable": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "text-moderation-007": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "text-moderation-latest": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "256-x-256/dall-e-2": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "512-x-512/dall-e-2": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "1024-x-1024/dall-e-2": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "hd/1024-x-1792/dall-e-3": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "hd/1792-x-1024/dall-e-3": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "hd/1024-x-1024/dall-e-3": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "standard/1024-x-1792/dall-e-3": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "standard/1792-x-1024/dall-e-3": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "standard/1024-x-1024/dall-e-3": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "whisper-1": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "tts-1": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "tts-1-hd": {"litellm_provider": "openai", "api_key_name": "OPENAI_API_KEY"}, "claude-instant-1": {"litellm_provider": "anthropic", "api_key_name": "ANTHROPIC_API_KEY"}, "groq/llama2-70b-4096": {"litellm_provider": "groq", "api_key_name": "GEMMA_API_KEY"}, "groq/llama3-8b-8192": {"litellm_provider": "groq", "api_key_name": "GEMMA_API_KEY"}, "groq/llama3-70b-8192": {"litellm_provider": "groq", "api_key_name": "GEMMA_API_KEY"}, "groq/llama-3.1-8b-instant": {"litellm_provider": "groq", "api_key_name": "GEMMA_API_KEY"}, "groq/llama-3.1-70b-versatile": {"litellm_provider": "groq", "api_key_name": "GEMMA_API_KEY"}, "groq/llama-3.1-405b-reasoning": {"litellm_provider": "groq", "api_key_name": "GEMMA_API_KEY"}, "groq/mixtral-8x7b-32768": {"litellm_provider": "groq", "api_key_name": "GEMMA_API_KEY"}, "groq/gemma-7b-it": {"litellm_provider": "groq", "api_key_name": "GEMMA_API_KEY"}, "groq/gemma2-9b-it": {"litellm_provider": "groq", "api_key_name": "GEMMA_API_KEY"}, "groq/llama3-groq-70b-8192-tool-use-preview": {"litellm_provider": "groq", "api_key_name": "GEMMA_API_KEY"}, "groq/llama3-groq-8b-8192-tool-use-preview": {"litellm_provider": "groq", "api_key_name": "GEMMA_API_KEY"}, "claude-instant-1.2": {"litellm_provider": "anthropic", "api_key_name": "ANTHROPIC_API_KEY"}, "claude-2": {"litellm_provider": "anthropic", "api_key_name": "ANTHROPIC_API_KEY"}, "claude-2.1": {"litellm_provider": "anthropic", "api_key_name": "ANTHROPIC_API_KEY"}, "claude-3-haiku-20240307": {"litellm_provider": "anthropic", "api_key_name": "ANTHROPIC_API_KEY"}, "claude-3-opus-20240229": {"litellm_provider": "anthropic", "api_key_name": "ANTHROPIC_API_KEY"}, "claude-3-sonnet-20240229": {"litellm_provider": "anthropic", "api_key_name": "ANTHROPIC_API_KEY"}, "claude-3-5-sonnet-20240620": {"litellm_provider": "anthropic", "api_key_name": "ANTHROPIC_API_KEY"}, "claude-3-5-sonnet-20241022": {"litellm_provider": "anthropic", "api_key_name": "ANTHROPIC_API_KEY"}, "gemini/gemini-1.5-flash-002": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-flash-001": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-flash": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-flash-latest": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-flash-8b-exp-0924": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-flash-exp-0827": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-flash-8b-exp-0827": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-pro": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-pro": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-pro-002": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-pro-001": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-pro-exp-0801": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-pro-exp-0827": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-1.5-pro-latest": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-pro-vision": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-gemma-2-27b-it": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "gemini/gemini-gemma-2-9b-it": {"litellm_provider": "gemini", "api_key_name": "GOOGLE_API_KEY"}, "anthropic/claude-3-5-sonnet-20241022": {"litellm_provider": "anthropic", "api_key_name": "ANTHROPIC_API_KEY"}}
35 changes: 29 additions & 6 deletions weave/trace_server/model_providers/model_providers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
import json
import logging
import os
from typing import Dict, TypedDict

import requests

model_providers_url = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
MODEL_PROVIDERS_FILE = "model_providers.json"
MODEL_PROVIDER_INFO_FILE = "model_providers.json"


PROVIDER_TO_API_KEY_NAME_MAP = {
"anthropic": "ANTHROPIC_API_KEY",
Expand All @@ -21,13 +23,27 @@ class LLMModelProviderInfo(TypedDict):
api_key_name: str


def fetch_model_to_provider_info_map(
cached_file_name: str = MODEL_PROVIDERS_FILE,
logger = logging.getLogger(__name__)


def read_model_to_provider_info_map(
file_name: str = MODEL_PROVIDER_INFO_FILE,
) -> Dict[str, LLMModelProviderInfo]:
full_path = os.path.join(os.path.dirname(__file__), cached_file_name)
if os.path.exists(full_path):
full_path = os.path.join(os.path.dirname(__file__), file_name)
try:
with open(full_path, "r") as f:
return json.load(f)
except Exception as e:
logger.error(
f"Failed to read model to provider info file at: {full_path}", exc_info=e
)
return {}


def main(
file_name: str = MODEL_PROVIDER_INFO_FILE,
) -> Dict[str, LLMModelProviderInfo]:
full_path = os.path.join(os.path.dirname(__file__), file_name)
try:
req = requests.get(model_providers_url)
req.raise_for_status()
Expand All @@ -43,7 +59,14 @@ def fetch_model_to_provider_info_map(
providers[k] = LLMModelProviderInfo(
litellm_provider=provider, api_key_name=api_key_name
)

os.makedirs(os.path.dirname(full_path), exist_ok=True)
with open(full_path, "w") as f:
json.dump(providers, f)
print(
f"Updated model to model provider info file at: {full_path}. {len(providers)} models updated."
)
return providers


if __name__ == "__main__":
main()
Loading