Skip to content

Commit 774966f

Browse files
lakshyaaglakshyaag
authored
Merge branch 'main' into feat/playground
2 parents 958b098 + db41d15 commit 774966f

10 files changed

Lines changed: 1980 additions & 1221 deletions

File tree

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ readme = "README.md"
99
requires-python = ">=3.11"
1010
dependencies = [
1111
"anthropic>=0.75.0",
12+
"google-genai>=1.56.0",
1213
"openai>=2.14.0",
1314
"portkey-ai>=2.1.0",
1415
"pytest>=9.0.2",

‎rlm/clients/__init__.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,15 @@ def get_client(
4444
from rlm.clients.anthropic import AnthropicClient
4545

4646
return AnthropicClient(**backend_kwargs)
47+
elif backend == "gemini":
48+
from rlm.clients.gemini import GeminiClient
49+
50+
return GeminiClient(**backend_kwargs)
51+
elif backend == "azure_openai":
52+
from rlm.clients.azure_openai import AzureOpenAIClient
53+
54+
return AzureOpenAIClient(**backend_kwargs)
4755
else:
4856
raise ValueError(
49-
f"Unknown backend: {backend}. Supported backends: ['openai', 'vllm', 'portkey', 'openrouter', 'litellm', 'anthropic']"
57+
f"Unknown backend: {backend}. Supported backends: ['openai', 'vllm', 'portkey', 'openrouter', 'litellm', 'anthropic', 'azure_openai', 'gemini']"
5058
)

‎rlm/clients/azure_openai.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
import os
2+
from collections import defaultdict
3+
from typing import Any
4+
5+
import openai
6+
from dotenv import load_dotenv
7+
8+
from rlm.clients.base_lm import BaseLM
9+
from rlm.core.types import ModelUsageSummary, UsageSummary
10+
11+
load_dotenv()
12+
13+
# Load API key from environment variable
14+
DEFAULT_AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
15+
16+
17+
class AzureOpenAIClient(BaseLM):
18+
"""
19+
LM Client for running models with the Azure OpenAI API.
20+
"""
21+
22+
def __init__(
23+
self,
24+
api_key: str | None = None,
25+
model_name: str | None = None,
26+
azure_endpoint: str | None = None,
27+
api_version: str | None = None,
28+
azure_deployment: str | None = None,
29+
**kwargs,
30+
):
31+
super().__init__(model_name=model_name, **kwargs)
32+
33+
if api_key is None:
34+
api_key = DEFAULT_AZURE_OPENAI_API_KEY
35+
36+
if azure_endpoint is None:
37+
azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
38+
39+
if api_version is None:
40+
api_version = os.getenv("AZURE_OPENAI_API_VERSION", "2024-02-01")
41+
42+
if azure_endpoint is None:
43+
raise ValueError(
44+
"azure_endpoint is required for Azure OpenAI client. "
45+
"Set it via argument or AZURE_OPENAI_ENDPOINT environment variable."
46+
)
47+
48+
self.client = openai.AzureOpenAI(
49+
api_key=api_key,
50+
azure_endpoint=azure_endpoint,
51+
api_version=api_version,
52+
azure_deployment=azure_deployment,
53+
)
54+
self.async_client = openai.AsyncAzureOpenAI(
55+
api_key=api_key,
56+
azure_endpoint=azure_endpoint,
57+
api_version=api_version,
58+
azure_deployment=azure_deployment,
59+
)
60+
self.model_name = model_name
61+
self.azure_deployment = azure_deployment
62+
63+
# Per-model usage tracking
64+
self.model_call_counts: dict[str, int] = defaultdict(int)
65+
self.model_input_tokens: dict[str, int] = defaultdict(int)
66+
self.model_output_tokens: dict[str, int] = defaultdict(int)
67+
self.model_total_tokens: dict[str, int] = defaultdict(int)
68+
69+
def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
70+
if isinstance(prompt, str):
71+
messages = [{"role": "user", "content": prompt}]
72+
elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
73+
messages = prompt
74+
else:
75+
raise ValueError(f"Invalid prompt type: {type(prompt)}")
76+
77+
model = model or self.model_name
78+
if not model:
79+
raise ValueError("Model name is required for Azure OpenAI client.")
80+
81+
response = self.client.chat.completions.create(
82+
model=model,
83+
messages=messages,
84+
)
85+
self._track_cost(response, model)
86+
return response.choices[0].message.content
87+
88+
async def acompletion(
89+
self, prompt: str | list[dict[str, Any]], model: str | None = None
90+
) -> str:
91+
if isinstance(prompt, str):
92+
messages = [{"role": "user", "content": prompt}]
93+
elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
94+
messages = prompt
95+
else:
96+
raise ValueError(f"Invalid prompt type: {type(prompt)}")
97+
98+
model = model or self.model_name
99+
if not model:
100+
raise ValueError("Model name is required for Azure OpenAI client.")
101+
102+
response = await self.async_client.chat.completions.create(
103+
model=model,
104+
messages=messages,
105+
)
106+
self._track_cost(response, model)
107+
return response.choices[0].message.content
108+
109+
def _track_cost(self, response: openai.ChatCompletion, model: str):
110+
self.model_call_counts[model] += 1
111+
112+
usage = getattr(response, "usage", None)
113+
if usage is None:
114+
raise ValueError("No usage data received. Tracking tokens not possible.")
115+
116+
self.model_input_tokens[model] += usage.prompt_tokens
117+
self.model_output_tokens[model] += usage.completion_tokens
118+
self.model_total_tokens[model] += usage.total_tokens
119+
120+
# Track last call for handler to read
121+
self.last_prompt_tokens = usage.prompt_tokens
122+
self.last_completion_tokens = usage.completion_tokens
123+
124+
def get_usage_summary(self) -> UsageSummary:
125+
model_summaries = {}
126+
for model in self.model_call_counts:
127+
model_summaries[model] = ModelUsageSummary(
128+
total_calls=self.model_call_counts[model],
129+
total_input_tokens=self.model_input_tokens[model],
130+
total_output_tokens=self.model_output_tokens[model],
131+
)
132+
return UsageSummary(model_usage_summaries=model_summaries)
133+
134+
def get_last_usage(self) -> ModelUsageSummary:
135+
return ModelUsageSummary(
136+
total_calls=1,
137+
total_input_tokens=self.last_prompt_tokens,
138+
total_output_tokens=self.last_completion_tokens,
139+
)

‎rlm/clients/gemini.py‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
import os
2+
from collections import defaultdict
3+
from typing import Any
4+
5+
from dotenv import load_dotenv
6+
from google import genai
7+
from google.genai import types
8+
9+
from rlm.clients.base_lm import BaseLM
10+
from rlm.core.types import ModelUsageSummary, UsageSummary
11+
12+
load_dotenv()
13+
14+
DEFAULT_GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
15+
16+
17+
class GeminiClient(BaseLM):
18+
"""
19+
LM Client for running models with the Google Gemini API.
20+
Uses the official google-genai SDK.
21+
"""
22+
23+
def __init__(
24+
self,
25+
api_key: str | None = None,
26+
model_name: str | None = "gemini-2.5-flash",
27+
**kwargs,
28+
):
29+
super().__init__(model_name=model_name, **kwargs)
30+
31+
if api_key is None:
32+
api_key = DEFAULT_GEMINI_API_KEY
33+
34+
if api_key is None:
35+
raise ValueError(
36+
"Gemini API key is required. Set GEMINI_API_KEY env var or pass api_key."
37+
)
38+
39+
self.client = genai.Client(api_key=api_key)
40+
self.model_name = model_name
41+
42+
# Per-model usage tracking
43+
self.model_call_counts: dict[str, int] = defaultdict(int)
44+
self.model_input_tokens: dict[str, int] = defaultdict(int)
45+
self.model_output_tokens: dict[str, int] = defaultdict(int)
46+
self.model_total_tokens: dict[str, int] = defaultdict(int)
47+
48+
# Last call tracking
49+
self.last_prompt_tokens = 0
50+
self.last_completion_tokens = 0
51+
52+
def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
53+
contents, system_instruction = self._prepare_contents(prompt)
54+
55+
model = model or self.model_name
56+
if not model:
57+
raise ValueError("Model name is required for Gemini client.")
58+
59+
config = None
60+
if system_instruction:
61+
config = types.GenerateContentConfig(system_instruction=system_instruction)
62+
63+
response = self.client.models.generate_content(
64+
model=model,
65+
contents=contents,
66+
config=config,
67+
)
68+
69+
self._track_cost(response, model)
70+
return response.text
71+
72+
async def acompletion(
73+
self, prompt: str | list[dict[str, Any]], model: str | None = None
74+
) -> str:
75+
contents, system_instruction = self._prepare_contents(prompt)
76+
77+
model = model or self.model_name
78+
if not model:
79+
raise ValueError("Model name is required for Gemini client.")
80+
81+
config = None
82+
if system_instruction:
83+
config = types.GenerateContentConfig(system_instruction=system_instruction)
84+
85+
# google-genai SDK supports async via aio interface
86+
response = await self.client.aio.models.generate_content(
87+
model=model,
88+
contents=contents,
89+
config=config,
90+
)
91+
92+
self._track_cost(response, model)
93+
return response.text
94+
95+
def _prepare_contents(
96+
self, prompt: str | list[dict[str, Any]]
97+
) -> tuple[list[types.Content] | str, str | None]:
98+
"""Prepare contents and extract system instruction for Gemini API."""
99+
system_instruction = None
100+
101+
if isinstance(prompt, str):
102+
return prompt, None
103+
104+
if isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
105+
# Convert OpenAI-style messages to Gemini format
106+
contents = []
107+
for msg in prompt:
108+
role = msg.get("role")
109+
content = msg.get("content", "")
110+
111+
if role == "system":
112+
# Gemini handles system instruction separately
113+
system_instruction = content
114+
elif role == "user":
115+
contents.append(types.Content(role="user", parts=[types.Part(text=content)]))
116+
elif role == "assistant":
117+
# Gemini uses "model" instead of "assistant"
118+
contents.append(types.Content(role="model", parts=[types.Part(text=content)]))
119+
else:
120+
# Default to user role for unknown roles
121+
contents.append(types.Content(role="user", parts=[types.Part(text=content)]))
122+
123+
return contents, system_instruction
124+
125+
raise ValueError(f"Invalid prompt type: {type(prompt)}")
126+
127+
def _track_cost(self, response: types.GenerateContentResponse, model: str):
128+
self.model_call_counts[model] += 1
129+
130+
# Extract token usage from response
131+
usage = response.usage_metadata
132+
if usage:
133+
input_tokens = usage.prompt_token_count or 0
134+
output_tokens = usage.candidates_token_count or 0
135+
136+
self.model_input_tokens[model] += input_tokens
137+
self.model_output_tokens[model] += output_tokens
138+
self.model_total_tokens[model] += input_tokens + output_tokens
139+
140+
# Track last call for handler to read
141+
self.last_prompt_tokens = input_tokens
142+
self.last_completion_tokens = output_tokens
143+
else:
144+
self.last_prompt_tokens = 0
145+
self.last_completion_tokens = 0
146+
147+
def get_usage_summary(self) -> UsageSummary:
148+
model_summaries = {}
149+
for model in self.model_call_counts:
150+
model_summaries[model] = ModelUsageSummary(
151+
total_calls=self.model_call_counts[model],
152+
total_input_tokens=self.model_input_tokens[model],
153+
total_output_tokens=self.model_output_tokens[model],
154+
)
155+
return UsageSummary(model_usage_summaries=model_summaries)
156+
157+
def get_last_usage(self) -> ModelUsageSummary:
158+
return ModelUsageSummary(
159+
total_calls=1,
160+
total_input_tokens=self.last_prompt_tokens,
161+
total_output_tokens=self.last_completion_tokens,
162+
)

‎rlm/core/rlm.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,7 @@ def completion(
186186
)
187187

188188
# Check if RLM is done and has a final answer.
189-
final_answer = find_final_answer(iteration.response)
189+
final_answer = find_final_answer(iteration.response, environment=environment)
190190
iteration.final_answer = final_answer
191191

192192
# If logger is used, log the iteration.
@@ -196,7 +196,7 @@ def completion(
196196
# Verbose output for this iteration
197197
self.verbose.print_iteration(iteration, i + 1)
198198

199-
if final_answer:
199+
if final_answer is not None:
200200
time_end = time.perf_counter()
201201
usage = lm_handler.get_usage_summary()
202202
self.verbose.print_final_answer(final_answer)

0 commit comments

Comments
 (0)