relia-sdk 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (47) hide show
  1. relia_sdk-0.1.0/.gitignore +70 -0
  2. relia_sdk-0.1.0/PKG-INFO +71 -0
  3. relia_sdk-0.1.0/pyproject.toml +79 -0
  4. relia_sdk-0.1.0/src/relia/__init__.py +29 -0
  5. relia_sdk-0.1.0/src/relia/__main__.py +4 -0
  6. relia_sdk-0.1.0/src/relia/agent/__init__.py +21 -0
  7. relia_sdk-0.1.0/src/relia/agent/auto.py +41 -0
  8. relia_sdk-0.1.0/src/relia/agent/core.py +121 -0
  9. relia_sdk-0.1.0/src/relia/agent/cost.py +60 -0
  10. relia_sdk-0.1.0/src/relia/agent/decorator.py +215 -0
  11. relia_sdk-0.1.0/src/relia/agent/interceptor.py +688 -0
  12. relia_sdk-0.1.0/src/relia/agent/wrappers/__init__.py +9 -0
  13. relia_sdk-0.1.0/src/relia/agent/wrappers/anthropic.py +186 -0
  14. relia_sdk-0.1.0/src/relia/agent/wrappers/langchain.py +154 -0
  15. relia_sdk-0.1.0/src/relia/agent/wrappers/openai.py +196 -0
  16. relia_sdk-0.1.0/src/relia/auto.py +22 -0
  17. relia_sdk-0.1.0/src/relia/cli/__init__.py +3 -0
  18. relia_sdk-0.1.0/src/relia/cli/api.py +149 -0
  19. relia_sdk-0.1.0/src/relia/cli/auth.py +44 -0
  20. relia_sdk-0.1.0/src/relia/cli/detect.py +68 -0
  21. relia_sdk-0.1.0/src/relia/cli/main.py +247 -0
  22. relia_sdk-0.1.0/src/relia/cli/patch.py +82 -0
  23. relia_sdk-0.1.0/src/relia/cli/prompts.py +153 -0
  24. relia_sdk-0.1.0/src/relia/config.py +106 -0
  25. relia_sdk-0.1.0/src/relia/core/__init__.py +18 -0
  26. relia_sdk-0.1.0/src/relia/core/sanitize.py +56 -0
  27. relia_sdk-0.1.0/src/relia/core/transport.py +102 -0
  28. relia_sdk-0.1.0/src/relia/core/types.py +37 -0
  29. relia_sdk-0.1.0/src/relia/integrations/__init__.py +11 -0
  30. relia_sdk-0.1.0/src/relia/integrations/django.py +53 -0
  31. relia_sdk-0.1.0/src/relia/integrations/fastapi.py +139 -0
  32. relia_sdk-0.1.0/src/relia/integrations/flask.py +73 -0
  33. relia_sdk-0.1.0/src/relia/py.typed +1 -0
  34. relia_sdk-0.1.0/src/relia/tracing/__init__.py +11 -0
  35. relia_sdk-0.1.0/src/relia/tracing/config.py +19 -0
  36. relia_sdk-0.1.0/src/relia/tracing/filters.py +46 -0
  37. relia_sdk-0.1.0/src/relia/tracing/otel.py +96 -0
  38. relia_sdk-0.1.0/src/relia/tracing/processor.py +41 -0
  39. relia_sdk-0.1.0/tests/test_agent_core.py +28 -0
  40. relia_sdk-0.1.0/tests/test_cli.py +17 -0
  41. relia_sdk-0.1.0/tests/test_cost.py +18 -0
  42. relia_sdk-0.1.0/tests/test_decorator.py +56 -0
  43. relia_sdk-0.1.0/tests/test_fastapi_middleware.py +28 -0
  44. relia_sdk-0.1.0/tests/test_interceptor.py +100 -0
  45. relia_sdk-0.1.0/tests/test_openai_wrapper.py +28 -0
  46. relia_sdk-0.1.0/tests/test_sanitize.py +35 -0
  47. relia_sdk-0.1.0/tests/test_tracing_otel.py +24 -0
@@ -0,0 +1,70 @@
1
+ # Dependencies
2
+ node_modules
3
+ node_modules/
4
+ **/node_modules
5
+ **/node_modules/
6
+ **/node_modules/**
7
+ .pnp
8
+ .pnp.*
9
+
10
+ # Python
11
+ __pycache__/
12
+ **/__pycache__/
13
+ *.py[cod]
14
+ *$py.class
15
+ .venv/
16
+ venv/
17
+ ENV/
18
+ .pytest_cache/
19
+ .uv/
20
+ .coverage
21
+ htmlcov/
22
+
23
+ # Build & Output
24
+ dist/
25
+ **/dist/
26
+ dist-ssr/
27
+ build/
28
+ out/
29
+ .next/
30
+ *.tsbuildinfo
31
+
32
+ # Environment & Secrets
33
+ .env
34
+ .env.*
35
+ !.env.example
36
+ !.env.template
37
+ *.pem
38
+ *.key
39
+
40
+ # Logs
41
+ logs/
42
+ *.log
43
+ npm-debug.log*
44
+ yarn-debug.log*
45
+ yarn-error.log*
46
+ pnpm-debug.log*
47
+ lerna-debug.log*
48
+ !src/components/logs/
49
+
50
+ # Databases
51
+ *.sqlite
52
+ *.sqlite3
53
+ *.db
54
+
55
+ # IDE & OS
56
+ .DS_Store
57
+ Thumbs.db
58
+ .idea/
59
+ .vscode/*
60
+ !.vscode/extensions.json
61
+ *.sw?
62
+ *.suo
63
+ *.ntvs*
64
+ *.njsproj
65
+ *.sln
66
+
67
+ # Other Tooling
68
+ .vercel/
69
+ .antigravitycli
70
+ .turbo/
@@ -0,0 +1,71 @@
1
+ Metadata-Version: 2.5
2
+ Name: relia-sdk
3
+ Version: 0.1.0
4
+ Summary: RELIA Observability & AI Agent Monitoring SDK for Python
5
+ Author-email: RELIA Team <support@tryrelia.com>
6
+ License: MIT
7
+ Keywords: ai-agents,apm,fastapi,llm-monitoring,observability,opentelemetry,relia,tracing
8
+ Classifier: Development Status :: 4 - Beta
9
+ Classifier: Intended Audience :: Developers
10
+ Classifier: License :: OSI Approved :: MIT License
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: Programming Language :: Python :: 3.9
13
+ Classifier: Programming Language :: Python :: 3.10
14
+ Classifier: Programming Language :: Python :: 3.11
15
+ Classifier: Programming Language :: Python :: 3.12
16
+ Classifier: Topic :: System :: Monitoring
17
+ Requires-Python: >=3.9
18
+ Requires-Dist: httpx>=0.24.0
19
+ Provides-Extra: agent
20
+ Requires-Dist: httpx>=0.24.0; extra == 'agent'
21
+ Provides-Extra: all
22
+ Requires-Dist: django>=4.0; extra == 'all'
23
+ Requires-Dist: fastapi>=0.100.0; extra == 'all'
24
+ Requires-Dist: flask>=2.0.0; extra == 'all'
25
+ Requires-Dist: httpx>=0.24.0; extra == 'all'
26
+ Requires-Dist: opentelemetry-api>=1.20.0; extra == 'all'
27
+ Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.20.0; extra == 'all'
28
+ Requires-Dist: opentelemetry-instrumentation-django>=0.40b0; extra == 'all'
29
+ Requires-Dist: opentelemetry-instrumentation-fastapi>=0.40b0; extra == 'all'
30
+ Requires-Dist: opentelemetry-instrumentation-flask>=0.40b0; extra == 'all'
31
+ Requires-Dist: opentelemetry-instrumentation-httpx>=0.40b0; extra == 'all'
32
+ Requires-Dist: opentelemetry-instrumentation-requests>=0.40b0; extra == 'all'
33
+ Requires-Dist: opentelemetry-sdk>=1.20.0; extra == 'all'
34
+ Provides-Extra: dev
35
+ Requires-Dist: fastapi>=0.100.0; extra == 'dev'
36
+ Requires-Dist: httpx>=0.24.0; extra == 'dev'
37
+ Requires-Dist: opentelemetry-api>=1.20.0; extra == 'dev'
38
+ Requires-Dist: opentelemetry-sdk>=1.20.0; extra == 'dev'
39
+ Requires-Dist: pytest-asyncio>=0.23.0; extra == 'dev'
40
+ Requires-Dist: pytest-mock>=3.12.0; extra == 'dev'
41
+ Requires-Dist: pytest>=8.0.0; extra == 'dev'
42
+ Provides-Extra: django
43
+ Requires-Dist: django>=4.0; extra == 'django'
44
+ Requires-Dist: opentelemetry-api>=1.20.0; extra == 'django'
45
+ Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.20.0; extra == 'django'
46
+ Requires-Dist: opentelemetry-instrumentation-django>=0.40b0; extra == 'django'
47
+ Requires-Dist: opentelemetry-instrumentation-httpx>=0.40b0; extra == 'django'
48
+ Requires-Dist: opentelemetry-instrumentation-requests>=0.40b0; extra == 'django'
49
+ Requires-Dist: opentelemetry-sdk>=1.20.0; extra == 'django'
50
+ Provides-Extra: fastapi
51
+ Requires-Dist: fastapi>=0.100.0; extra == 'fastapi'
52
+ Requires-Dist: opentelemetry-api>=1.20.0; extra == 'fastapi'
53
+ Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.20.0; extra == 'fastapi'
54
+ Requires-Dist: opentelemetry-instrumentation-fastapi>=0.40b0; extra == 'fastapi'
55
+ Requires-Dist: opentelemetry-instrumentation-httpx>=0.40b0; extra == 'fastapi'
56
+ Requires-Dist: opentelemetry-instrumentation-requests>=0.40b0; extra == 'fastapi'
57
+ Requires-Dist: opentelemetry-sdk>=1.20.0; extra == 'fastapi'
58
+ Provides-Extra: flask
59
+ Requires-Dist: flask>=2.0.0; extra == 'flask'
60
+ Requires-Dist: opentelemetry-api>=1.20.0; extra == 'flask'
61
+ Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.20.0; extra == 'flask'
62
+ Requires-Dist: opentelemetry-instrumentation-flask>=0.40b0; extra == 'flask'
63
+ Requires-Dist: opentelemetry-instrumentation-httpx>=0.40b0; extra == 'flask'
64
+ Requires-Dist: opentelemetry-instrumentation-requests>=0.40b0; extra == 'flask'
65
+ Requires-Dist: opentelemetry-sdk>=1.20.0; extra == 'flask'
66
+ Provides-Extra: otel
67
+ Requires-Dist: opentelemetry-api>=1.20.0; extra == 'otel'
68
+ Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.20.0; extra == 'otel'
69
+ Requires-Dist: opentelemetry-instrumentation-httpx>=0.40b0; extra == 'otel'
70
+ Requires-Dist: opentelemetry-instrumentation-requests>=0.40b0; extra == 'otel'
71
+ Requires-Dist: opentelemetry-sdk>=1.20.0; extra == 'otel'
@@ -0,0 +1,79 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "relia-sdk"
7
+ version = "0.1.0"
8
+ description = "RELIA Observability & AI Agent Monitoring SDK for Python"
9
+ requires-python = ">=3.9"
10
+ license = { text = "MIT" }
11
+ authors = [
12
+ { name = "RELIA Team", email = "support@tryrelia.com" }
13
+ ]
14
+ keywords = ["observability", "opentelemetry", "tracing", "apm", "ai-agents", "llm-monitoring", "fastapi", "relia"]
15
+ classifiers = [
16
+ "Development Status :: 4 - Beta",
17
+ "Intended Audience :: Developers",
18
+ "License :: OSI Approved :: MIT License",
19
+ "Programming Language :: Python :: 3",
20
+ "Programming Language :: Python :: 3.9",
21
+ "Programming Language :: Python :: 3.10",
22
+ "Programming Language :: Python :: 3.11",
23
+ "Programming Language :: Python :: 3.12",
24
+ "Topic :: System :: Monitoring",
25
+ ]
26
+ dependencies = [
27
+ "httpx>=0.24.0",
28
+ ]
29
+
30
+ [project.optional-dependencies]
31
+ otel = [
32
+ "opentelemetry-api>=1.20.0",
33
+ "opentelemetry-sdk>=1.20.0",
34
+ "opentelemetry-exporter-otlp-proto-http>=1.20.0",
35
+ "opentelemetry-instrumentation-requests>=0.40b0",
36
+ "opentelemetry-instrumentation-httpx>=0.40b0",
37
+ ]
38
+ fastapi = [
39
+ "relia-sdk[otel]",
40
+ "opentelemetry-instrumentation-fastapi>=0.40b0",
41
+ "fastapi>=0.100.0",
42
+ ]
43
+ flask = [
44
+ "relia-sdk[otel]",
45
+ "opentelemetry-instrumentation-flask>=0.40b0",
46
+ "flask>=2.0.0",
47
+ ]
48
+ django = [
49
+ "relia-sdk[otel]",
50
+ "opentelemetry-instrumentation-django>=0.40b0",
51
+ "django>=4.0",
52
+ ]
53
+ agent = [
54
+ "httpx>=0.24.0",
55
+ ]
56
+ all = [
57
+ "relia-sdk[fastapi,flask,django,agent]",
58
+ ]
59
+ dev = [
60
+ "pytest>=8.0.0",
61
+ "pytest-asyncio>=0.23.0",
62
+ "pytest-mock>=3.12.0",
63
+ "fastapi>=0.100.0",
64
+ "httpx>=0.24.0",
65
+ "opentelemetry-api>=1.20.0",
66
+ "opentelemetry-sdk>=1.20.0",
67
+ ]
68
+
69
+ [project.scripts]
70
+ relia = "relia.cli.main:main"
71
+
72
+ [tool.hatch.build.targets.wheel]
73
+ packages = ["src/relia"]
74
+
75
+ [tool.pytest.ini_options]
76
+ testpaths = ["tests"]
77
+ pythonpath = ["src"]
78
+ asyncio_mode = "auto"
79
+
@@ -0,0 +1,29 @@
1
+ from relia.agent import estimate_cost, report_run, tool, trace, trace_run
2
+ from relia.agent.auto import init as init_agent
3
+ from relia.config import ReliaConfig
4
+ from relia.core.sanitize import sanitize_data, sanitize_object, sanitize_payload_string
5
+ from relia.core.transport import flush
6
+ from relia.integrations.fastapi import instrument_fastapi
7
+ from relia.tracing import start
8
+
9
+ __version__ = "0.1.0"
10
+
11
+ init = init_agent
12
+
13
+ __all__ = [
14
+ "start",
15
+ "init",
16
+ "instrument_fastapi",
17
+ "trace",
18
+ "tool",
19
+ "trace_run",
20
+ "report_run",
21
+ "estimate_cost",
22
+ "flush",
23
+ "sanitize_data",
24
+ "sanitize_object",
25
+ "sanitize_payload_string",
26
+ "ReliaConfig",
27
+ "__version__",
28
+ ]
29
+
@@ -0,0 +1,4 @@
1
+ from relia.cli.main import main
2
+
3
+ if __name__ == "__main__":
4
+ main()
@@ -0,0 +1,21 @@
1
+ from relia.agent.auto import init
2
+ from relia.agent.core import extract_prompt_from_messages, report_run, resolve_agent_id
3
+ from relia.agent.cost import estimate_cost
4
+ from relia.agent.decorator import tool, trace, trace_run
5
+ from relia.agent.wrappers import ReliaCallbackHandler, patch_anthropic, patch_openai
6
+ from relia.core.transport import flush
7
+
8
+ __all__ = [
9
+ "init",
10
+ "trace",
11
+ "tool",
12
+ "trace_run",
13
+ "report_run",
14
+ "estimate_cost",
15
+ "resolve_agent_id",
16
+ "extract_prompt_from_messages",
17
+ "patch_openai",
18
+ "patch_anthropic",
19
+ "ReliaCallbackHandler",
20
+ "flush",
21
+ ]
@@ -0,0 +1,41 @@
1
+ import logging
2
+ from typing import Optional
3
+ from relia.agent.wrappers.anthropic import patch_anthropic
4
+ from relia.agent.wrappers.openai import patch_openai
5
+ from relia.config import ReliaConfig
6
+
7
+ logger = logging.getLogger("relia")
8
+
9
+ _initialized = False
10
+
11
+ def init(
12
+ agent_id: Optional[str] = None,
13
+ agent_name: Optional[str] = None,
14
+ endpoint: Optional[str] = None,
15
+ ) -> bool:
16
+ """
17
+ Auto-initialize AI agent tracking by patching installed LLM libraries (OpenAI, Anthropic).
18
+ """
19
+ global _initialized
20
+ if _initialized:
21
+ return False
22
+
23
+ config = ReliaConfig.from_env(
24
+ agent_id=agent_id,
25
+ agent_name=agent_name,
26
+ agent_endpoint=endpoint,
27
+ )
28
+
29
+ from relia.agent.interceptor import init_interceptor
30
+
31
+ interceptor_success = init_interceptor()
32
+ openai_success = patch_openai()
33
+ anthropic_success = patch_anthropic()
34
+
35
+ _initialized = True
36
+
37
+ if interceptor_success or openai_success or anthropic_success:
38
+ logger.info("[relia] agent tracing initialized for %s", config.agent_name)
39
+ return True
40
+
41
+ return False
@@ -0,0 +1,121 @@
1
+ import hashlib
2
+ import os
3
+ from typing import Any, Dict, List, Optional
4
+ from relia.agent.cost import estimate_cost
5
+ from relia.config import DEFAULT_AGENT_ENDPOINT, ReliaConfig
6
+ from relia.core.transport import report
7
+ from relia.core.types import AgentRunPayload, AgentStep
8
+
9
+ def resolve_agent_id(agent_name: str) -> str:
10
+ env_id = os.getenv("RELIA_AGENT_ID")
11
+ if env_id:
12
+ return env_id
13
+ base = agent_name or "agent"
14
+ return hashlib.sha1(base.encode("utf-8")).hexdigest()[:24]
15
+
16
+ def extract_prompt_from_messages(messages: Any) -> str:
17
+ if isinstance(messages, str):
18
+ return messages[:5000]
19
+ if not isinstance(messages, list) or not messages:
20
+ return ""
21
+
22
+ # Find the last user message, fallback to last message
23
+ user_msg = None
24
+ for m in reversed(messages):
25
+ role = getattr(m, "role", None) or (m.get("role") if isinstance(m, dict) else None)
26
+ if role == "user":
27
+ user_msg = m
28
+ break
29
+ if not user_msg:
30
+ user_msg = messages[-1]
31
+
32
+ content = getattr(user_msg, "content", None) or (user_msg.get("content") if isinstance(user_msg, dict) else user_msg)
33
+ if isinstance(content, str):
34
+ return content[:5000]
35
+ if isinstance(content, list):
36
+ text_parts = []
37
+ for p in content:
38
+ if isinstance(p, str):
39
+ text_parts.append(p)
40
+ elif isinstance(p, dict):
41
+ text_parts.append(str(p.get("text", "")))
42
+ elif hasattr(p, "text"):
43
+ text_parts.append(str(getattr(p, "text", "")))
44
+ return " ".join(filter(None, text_parts))[:5000]
45
+
46
+ return str(content or "")[:5000]
47
+
48
+ def report_run(
49
+ agent_id: Optional[str] = None,
50
+ name: Optional[str] = None,
51
+ service_name: Optional[str] = None,
52
+ primary_model: str = "ai-model",
53
+ prompt: str = "",
54
+ answer: str = "",
55
+ duration_ms: int = 0,
56
+ tokens_in: int = 0,
57
+ tokens_out: int = 0,
58
+ status: str = "ok",
59
+ steps: Optional[List[Dict[str, Any]]] = None,
60
+ endpoint: Optional[str] = None,
61
+ user: Optional[str] = None,
62
+ error: Optional[Any] = None,
63
+ cost_usd: Optional[float] = None,
64
+ version: str = "1.0.0",
65
+ ) -> None:
66
+ """
67
+ Report an AI Agent execution run directly to the RELIA backend.
68
+ """
69
+ config = ReliaConfig.from_env()
70
+ final_name = name or service_name or config.agent_name or "agent"
71
+ final_agent_id = agent_id or config.agent_id or final_name
72
+ final_endpoint = endpoint or config.agent_endpoint or DEFAULT_AGENT_ENDPOINT
73
+ final_user = user or config.agent_user or ""
74
+
75
+
76
+ if cost_usd is None:
77
+ cost_usd = estimate_cost(primary_model, tokens_in, tokens_out)
78
+
79
+ final_status = status
80
+ if error is not None:
81
+ final_status = "error"
82
+
83
+ raw_steps = steps or []
84
+ formatted_steps: List[AgentStep] = []
85
+ for s in raw_steps:
86
+ formatted_steps.append({
87
+ "type": s.get("type", "model"),
88
+ "name": str(s.get("name", "step"))[:100],
89
+ "detail": str(s.get("detail", ""))[:500] if s.get("detail") is not None else "",
90
+ "offset_ms": int(s.get("offset_ms", 0)),
91
+ "duration_ms": int(s.get("duration_ms", 0)),
92
+ "tokens": s.get("tokens"),
93
+ "tokens_in": s.get("tokens_in"),
94
+ "tokens_out": s.get("tokens_out"),
95
+ "provider": str(s.get("provider", "")) if s.get("provider") else None,
96
+ "model": str(s.get("model", "")) if s.get("model") else None,
97
+ "status": "error" if s.get("status") == "error" else "ok",
98
+ "input": str(s.get("input", ""))[:10000] if s.get("input") is not None else None,
99
+ "output": str(s.get("output", ""))[:10000] if s.get("output") is not None else None,
100
+ })
101
+
102
+
103
+ payload: AgentRunPayload = {
104
+ "agent_id": final_agent_id,
105
+ "name": final_name,
106
+ "service_name": final_name,
107
+ "primary_model": primary_model or "ai-model",
108
+ "version": version,
109
+ "endpoint": final_endpoint.replace("/runs", ""),
110
+ "user": final_user,
111
+ "duration_ms": max(0, duration_ms),
112
+ "tokens_in": max(0, tokens_in),
113
+ "tokens_out": max(0, tokens_out),
114
+ "cost_usd": cost_usd,
115
+ "status": "error" if final_status == "error" else "ok",
116
+ "prompt": str(prompt)[:5000],
117
+ "answer": str(answer)[:5000],
118
+ "steps": formatted_steps,
119
+ }
120
+
121
+ report(payload, final_endpoint)
@@ -0,0 +1,60 @@
1
+ import re
2
+ from typing import Dict, Tuple
3
+
4
+ PRICES_PER_1M: Dict[str, Tuple[float, float]] = {
5
+ # model: (input_cost_per_1M, output_cost_per_1M)
6
+ "gpt-5.5-instant": (0.5, 2.0),
7
+ "gpt-5": (1.25, 10.0),
8
+ "gpt-4o": (2.5, 10.0),
9
+ "gpt-4o-mini": (0.15, 0.6),
10
+ "gpt-4.1": (2.0, 8.0),
11
+ "gpt-4.1-mini": (0.4, 1.6),
12
+ "gpt-4-turbo": (10.0, 30.0),
13
+ "o1": (15.0, 60.0),
14
+ "o1-mini": (1.1, 4.4),
15
+ "o3": (2.0, 8.0),
16
+ "o3-mini": (1.1, 4.4),
17
+ "o4-mini": (1.1, 4.4),
18
+ "claude-opus": (15.0, 75.0),
19
+ "claude-opus-4": (15.0, 75.0),
20
+ "claude-sonnet": (3.0, 15.0),
21
+ "claude-sonnet-4": (3.0, 15.0),
22
+ "claude-3-7-sonnet": (3.0, 15.0),
23
+ "claude-3-5-sonnet": (3.0, 15.0),
24
+ "claude-3-haiku": (0.25, 1.25),
25
+ "claude-3-5-haiku": (0.8, 4.0),
26
+ "deepseek-chat": (0.14, 0.28),
27
+ "deepseek-reasoner": (0.55, 2.19),
28
+ "gemini-2.5-pro": (1.25, 10.0),
29
+ "gemini-2.5-flash": (0.3, 2.5),
30
+ "gemini-2.0-flash": (0.1, 0.4),
31
+ "gemini-1.5-pro": (1.25, 5.0),
32
+ "gemini-1.5-flash": (0.075, 0.3),
33
+ }
34
+
35
+ def normalize_model(model: str) -> str:
36
+ cleaned = model.lower()
37
+ cleaned = re.sub(r"^[a-z0-9-]+\/", "", cleaned)
38
+ cleaned = re.sub(r":.*$", "", cleaned)
39
+ return cleaned.strip()
40
+
41
+ def estimate_cost(model: str, tokens_in: int, tokens_out: int) -> float:
42
+ """
43
+ Estimate USD cost based on tokens_in and tokens_out for the given model ID.
44
+ """
45
+ if not model or (tokens_in <= 0 and tokens_out <= 0):
46
+ return 0.0
47
+
48
+ key = normalize_model(model)
49
+ price = (
50
+ PRICES_PER_1M.get(key)
51
+ or PRICES_PER_1M.get(re.sub(r"-\d+$", "", key))
52
+ or PRICES_PER_1M.get(key.split("-")[0])
53
+ )
54
+
55
+ if not price:
56
+ return 0.0
57
+
58
+ in_price, out_price = price
59
+ total = (tokens_in / 1_000_000.0) * in_price + (tokens_out / 1_000_000.0) * out_price
60
+ return round(total, 6)