Skip to content

Commit c344bef

Browse files
fix: Add base_url and api_key parameters to Agent class for Ollama support
This allows users to connect to remote Ollama instances without relying on environment variables. Usage: ```python agent = Agent( llm="ollama/llama3.2", base_url="http://remote-host:11434" ) ``` Fixes #482 Co-authored-by: Mervin Praison <[email protected]>
1 parent fe29aea commit c344bef

File tree

2 files changed

+76
-3
lines changed

2 files changed

+76
-3
lines changed
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# Example: Using Ollama with a remote host via base_url parameter
2+
# pip install praisonaiagents
3+
4+
from praisonaiagents import Agent
5+
6+
# Method 1: Using base_url parameter directly (NEW)
7+
agent = Agent(
8+
instructions="You are a helpful assistant",
9+
llm="ollama/llama3.2",
10+
base_url="http://remote-host:11434"
11+
)
12+
13+
response = agent.start("Why is the sky blue?")
14+
print(response)
15+
16+
# Method 2: Using dict configuration
17+
agent2 = Agent(
18+
instructions="You are a helpful assistant",
19+
llm={
20+
"model": "ollama/llama3.2",
21+
"base_url": "http://remote-host:11434"
22+
}
23+
)
24+
25+
response2 = agent2.start("What is the capital of France?")
26+
print(response2)
27+
28+
# Method 3: Using environment variable (existing method)
29+
# export OPENAI_BASE_URL=http://remote-host:11434
30+
# Then just:
31+
# agent3 = Agent(
32+
# instructions="You are a helpful assistant",
33+
# llm="ollama/llama3.2"
34+
# )

src/praisonai-agents/praisonaiagents/agent/agent.py

Lines changed: 42 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -373,7 +373,9 @@ def __init__(
373373
user_id: Optional[str] = None,
374374
reasoning_steps: bool = False,
375375
guardrail: Optional[Union[Callable[['TaskOutput'], Tuple[bool, Any]], str]] = None,
376-
max_guardrail_retries: int = 3
376+
max_guardrail_retries: int = 3,
377+
base_url: Optional[str] = None,
378+
api_key: Optional[str] = None
377379
):
378380
"""Initialize an Agent instance.
379381
@@ -457,6 +459,10 @@ def __init__(
457459
description string for LLM-based validation. Defaults to None.
458460
max_guardrail_retries (int, optional): Maximum number of retry attempts when guardrail
459461
validation fails before giving up. Defaults to 3.
462+
base_url (Optional[str], optional): Base URL for custom LLM endpoints (e.g., Ollama).
463+
If provided, automatically creates a custom LLM instance. Defaults to None.
464+
api_key (Optional[str], optional): API key for LLM provider. If not provided,
465+
falls back to environment variables. Defaults to None.
460466
461467
Raises:
462468
ValueError: If all of name, role, goal, backstory, and instructions are None.
@@ -503,10 +509,40 @@ def __init__(
503509
# Check for model name in environment variable if not provided
504510
self._using_custom_llm = False
505511

512+
# If base_url is provided, always create a custom LLM instance
513+
if base_url:
514+
try:
515+
from ..llm.llm import LLM
516+
# Handle different llm parameter types with base_url
517+
if isinstance(llm, dict):
518+
# Merge base_url and api_key into the dict
519+
llm_config = llm.copy()
520+
llm_config['base_url'] = base_url
521+
if api_key:
522+
llm_config['api_key'] = api_key
523+
self.llm_instance = LLM(**llm_config)
524+
else:
525+
# Create LLM with model string and base_url
526+
model_name = llm or os.getenv('OPENAI_MODEL_NAME', 'gpt-4o')
527+
self.llm_instance = LLM(
528+
model=model_name,
529+
base_url=base_url,
530+
api_key=api_key
531+
)
532+
self._using_custom_llm = True
533+
except ImportError as e:
534+
raise ImportError(
535+
"LLM features requested but dependencies not installed. "
536+
"Please install with: pip install \"praisonaiagents[llm]\""
537+
) from e
506538
# If the user passes a dictionary (for advanced configuration)
507-
if isinstance(llm, dict) and "model" in llm:
539+
elif isinstance(llm, dict) and "model" in llm:
508540
try:
509541
from ..llm.llm import LLM
542+
# Add api_key if provided and not in dict
543+
if api_key and 'api_key' not in llm:
544+
llm = llm.copy()
545+
llm['api_key'] = api_key
510546
self.llm_instance = LLM(**llm) # Pass all dict items as kwargs
511547
self._using_custom_llm = True
512548
except ImportError as e:
@@ -519,7 +555,10 @@ def __init__(
519555
try:
520556
from ..llm.llm import LLM
521557
# Pass the entire string so LiteLLM can parse provider/model
522-
self.llm_instance = LLM(model=llm)
558+
llm_params = {'model': llm}
559+
if api_key:
560+
llm_params['api_key'] = api_key
561+
self.llm_instance = LLM(**llm_params)
523562
self._using_custom_llm = True
524563

525564
# Ensure tools are properly accessible when using custom LLM

0 commit comments

Comments
 (0)