mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-08 01:08:39 -10:00
Compare commits
99
Commits
No files matched your search
@@ -158,6 +158,8 @@ jobs:
|
||||
tests/unit/test_adapter_registry.py tests/unit/test_token_budget.py
|
||||
tests/unit/test_streaming_events.py tests/unit/test_tool_decorator_approval.py
|
||||
tests/unit/test_task_execution_config.py
|
||||
tests/unit/test_agent_unknown_kwarg_guidance.py
|
||||
tests/unit/test_agent_instantiation_lazy.py tests/unit/test_agent_os_exports.py
|
||||
# Both tests previously deselected here now pass in combination.
|
||||
# test_lazy_litellm_import was looking the loader module up by
|
||||
# name and could get a different object than the one get_litellm()
|
||||
|
||||
@@ -70,7 +70,7 @@ print(agent.start("In two sentences, why do some companies keep their data on th
|
||||
|
||||
Run the same agent on a model on your own computer with [Ollama](https://ollama.com). You need no API key, and your prompts and answers stay on the machine.
|
||||
|
||||
Install the package and pull the model while you still have internet; after that, the agent runs with no network at all.
|
||||
Install the package and pull the model while you still have internet; after that, the agent runs with no network at all. If you installed only the core library, use `pip install "praisonaiagents[llm]"`, which adds the local-model support.
|
||||
|
||||
```bash
|
||||
# Start Ollama first (the app, or "ollama serve" in another terminal)
|
||||
@@ -96,11 +96,13 @@ vLLM and LM Studio work the same way: use `hosted_vllm/` (vLLM) or `lm_studio/`
|
||||
## How the handover works (early access)
|
||||
|
||||
1. **Watch:** records how each task is done today.
|
||||
2. **Test:** turns that real work into tests.
|
||||
2. **Test:** your own AI takes an exam made from real past work it never studied.
|
||||
3. **Hand over:** a task moves to your own AI only when it passes.
|
||||
4. **Roll back:** if quality drops, the task moves back automatically.
|
||||
5. **Record:** every handover leaves a record your regulator can check.
|
||||
|
||||
**What is "your own AI"?** An open model, such as Llama, Qwen or Gemma, running on your own computers. Use it as it is, or train it on your own past work with [PraisonAI Train](https://praison.ai/docs/train) on your own GPU servers. Your team decides when to train; nothing trains automatically. Either way, a task moves to it only after it passes that exam.
|
||||
|
||||
Highly regulated teams, such as banks, insurers and hospitals, want AI but can't send customer data to outside models. From 6 January 2027, OpenAI also takes no new self-serve fine-tuning jobs ([notice](https://developers.openai.com/api/docs/deprecations)), so teams that customise models need models they own. The handover lets these teams move work to AI they own, one task at a time, with a record that shows why each step was trusted.
|
||||
|
||||
It is built on parts of this framework that are live today: tracing, evaluation, fine-tuning, model routing and local models. The handover itself is in early access, and we are taking on our first design partners. To take part, see [praison.ai/enterprise](https://praison.ai/enterprise/).
|
||||
|
||||
@@ -132,17 +132,17 @@ Methods:
|
||||
* <code title="class DeepResearchAgent">DeepResearchAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/deep_research_agent.py">openai_client</a>()</code>
|
||||
* <code title="class DeepResearchAgent">DeepResearchAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/deep_research_agent.py">research</a>(query: str, instructions: Optional[str] = None, model: Optional[str] = None, summary_mode: Optional[Literal['auto', 'detailed', 'concise']] = None, web_search: Optional[bool] = None, code_interpreter: Optional[bool] = None, mcp_servers: Optional[List[Dict[str, Any]]] = None, file_ids: Optional[List[str]] = None, file_search: Optional[bool] = None, file_search_stores: Optional[List[str]] = None, stream: bool = True) -> DeepResearchResponse</code>
|
||||
* <code title="class DeepResearchAgent">DeepResearchAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/deep_research_agent.py">rewrite_query</a>(query: str, model: Optional[str] = None) -> str</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">achat</a>(prompt: str, temperature: float = 0.2, tools: Optional[List[Any]] = None, output_json: Optional[str] = None, output_pydantic: Optional[Any] = None, reasoning_steps: bool = False, **kwargs) -> Union[str, Dict[str, Any]]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">aedit</a>(image: str, prompt: str, mask: Optional[str] = None, n: int = 1, size: Optional[str] = None, **kwargs) -> Dict[str, Any]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">agenerate</a>(prompt: str, **kwargs) -> Dict[str, Any]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">agenerate_image</a>(prompt: str, **kwargs) -> Dict[str, Any]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">avariation</a>(image: str, n: int = 1, size: Optional[str] = None, **kwargs) -> Dict[str, Any]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">chat</a>(prompt: str, **kwargs) -> Dict[str, Any]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">edit</a>(image: str, prompt: str, mask: Optional[str] = None, n: int = 1, size: Optional[str] = None, **kwargs) -> Dict[str, Any]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">generate</a>(prompt: str, **kwargs) -> Dict[str, Any]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">generate_image</a>(prompt: str, **kwargs) -> Dict[str, Any]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">achat</a>(prompt: str, temperature: float = 0.2, tools: Optional[List[Any]] = None, output_json: Optional[str] = None, output_pydantic: Optional[Any] = None, reasoning_steps: bool = False, **kwargs) -> Union[ImageResult, Dict[str, Any]]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">aedit</a>(image: str, prompt: str, mask: Optional[str] = None, n: int = 1, size: Optional[str] = None, **kwargs) -> ImageResult</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">agenerate</a>(prompt: str, **kwargs) -> ImageResult</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">agenerate_image</a>(prompt: str, **kwargs) -> ImageResult</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">avariation</a>(image: str, n: int = 1, size: Optional[str] = None, **kwargs) -> ImageResult</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">chat</a>(prompt: str, **kwargs) -> Union[ImageResult, Dict[str, Any]]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">edit</a>(image: str, prompt: str, mask: Optional[str] = None, n: int = 1, size: Optional[str] = None, **kwargs) -> ImageResult</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">generate</a>(prompt: str, **kwargs) -> ImageResult</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">generate_image</a>(prompt: str, **kwargs) -> ImageResult</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">litellm</a>()</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">variation</a>(image: str, n: int = 1, size: Optional[str] = None, **kwargs) -> Dict[str, Any]</code>
|
||||
* <code title="class ImageAgent">ImageAgent.<a href="./src/praisonai-agents/praisonaiagents/agent/image_agent.py">variation</a>(image: str, n: int = 1, size: Optional[str] = None, **kwargs) -> ImageResult</code>
|
||||
* <code title="class PlanningAgent">PlanningAgent.<a href="./src/praisonai-agents/praisonaiagents/planning/planner.py">analyze_context</a>(context: str) -> str</code>
|
||||
* <code title="class PlanningAgent">PlanningAgent.<a href="./src/praisonai-agents/praisonaiagents/planning/planner.py">analyze_context_sync</a>(context: str) -> str</code>
|
||||
* <code title="class PlanningAgent">PlanningAgent.<a href="./src/praisonai-agents/praisonaiagents/planning/planner.py">create_plan</a>(request: str, agents: List['Agent'], tasks: Optional[List['Task']] = None, context: Optional[str] = None) -> Plan</code>
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.23" \
|
||||
"praisonai>=4.7.24" \
|
||||
"praisonai[chat]" \
|
||||
"embedchain[github,youtube]"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.23" \
|
||||
"praisonai>=4.7.24" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.23" \
|
||||
"praisonai>=4.7.24" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -5,23 +5,39 @@ PraisonAI Agents - Real Execution Benchmark
|
||||
Benchmarks ACTUAL agent execution with LLM API calls.
|
||||
Tests single agent running a simple task across frameworks.
|
||||
|
||||
Each framework gets untimed warm-up calls first (imports, client set-up and
|
||||
the first connection are not what this measures), then the timed runs are
|
||||
interleaved round by round so network drift hits every framework alike.
|
||||
A run that raises or returns an empty answer is counted as a failure, not
|
||||
timed. The summary reports the median.
|
||||
|
||||
Usage:
|
||||
export OPENAI_API_KEY=your_key
|
||||
python benchmarks/execution_benchmark.py
|
||||
python benchmarks/execution_benchmark.py --no-save
|
||||
python benchmarks/execution_benchmark.py --iterations 5
|
||||
python benchmarks/execution_benchmark.py --save
|
||||
python benchmarks/execution_benchmark.py --iterations 10
|
||||
python benchmarks/execution_benchmark.py --model gpt-4o --iterations 10
|
||||
"""
|
||||
|
||||
import time
|
||||
import os
|
||||
import argparse
|
||||
import os
|
||||
import statistics
|
||||
import time
|
||||
from importlib.metadata import version as get_version
|
||||
|
||||
# Defaults (can be overridden via CLI)
|
||||
DEFAULT_MODEL = "gpt-4o-mini"
|
||||
DEFAULT_ITERATIONS = 3
|
||||
DEFAULT_ITERATIONS = 10
|
||||
DEFAULT_WARMUP = 1
|
||||
DEFAULT_PROMPT = "What is 2+2? Reply with just the number."
|
||||
INSTRUCTIONS = "You are a helpful assistant. Be very brief."
|
||||
|
||||
METHODS = {
|
||||
'PraisonAI': 'agent.start()',
|
||||
'PraisonAI (LiteLLM)': 'agent.start()',
|
||||
'Agno': 'agent.run()',
|
||||
'CrewAI': 'crew.kickoff()',
|
||||
}
|
||||
|
||||
|
||||
def get_package_versions():
|
||||
@@ -35,260 +51,235 @@ def get_package_versions():
|
||||
for name, pkg in packages.items():
|
||||
try:
|
||||
versions[name] = get_version(pkg)
|
||||
except:
|
||||
except Exception:
|
||||
versions[name] = 'not installed'
|
||||
return versions
|
||||
|
||||
|
||||
def benchmark_praisonai(model, iterations, prompt):
|
||||
"""Benchmark PraisonAI agent execution using agent.start()"""
|
||||
def praisonai_runner(model, prompt):
|
||||
"""PraisonAI: a new agent per run; only agent.start() is timed."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
print("\n--- PraisonAI ---")
|
||||
print("Method: agent.start()")
|
||||
|
||||
times = []
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
name="Calculator",
|
||||
instructions="You are a helpful assistant. Be very brief.",
|
||||
llm=model
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
response = agent.start(prompt, output="silent")
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
times.append(elapsed)
|
||||
print(f" Run {i+1}: {elapsed:.2f}s - Response: {str(response)[:50]}")
|
||||
|
||||
avg = sum(times) / len(times)
|
||||
print(f" Average: {avg:.2f}s")
|
||||
return avg
|
||||
|
||||
def build():
|
||||
agent = Agent(name="Calculator", instructions=INSTRUCTIONS, llm=model)
|
||||
return lambda: agent.start(prompt, output="silent")
|
||||
return build
|
||||
|
||||
|
||||
def benchmark_praisonai_litellm(model, iterations, prompt):
|
||||
"""Benchmark PraisonAI with LiteLLM backend using agent.start()"""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
print("\n--- PraisonAI (LiteLLM) ---")
|
||||
print("Method: agent.start()")
|
||||
|
||||
times = []
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
name="Calculator",
|
||||
instructions="You are a helpful assistant. Be very brief.",
|
||||
llm=f"openai/{model}"
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
response = agent.start(prompt, output="silent")
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
times.append(elapsed)
|
||||
print(f" Run {i+1}: {elapsed:.2f}s - Response: {str(response)[:50]}")
|
||||
|
||||
avg = sum(times) / len(times)
|
||||
print(f" Average: {avg:.2f}s")
|
||||
return avg
|
||||
|
||||
|
||||
def benchmark_agno(model, iterations, prompt):
|
||||
"""Benchmark Agno agent execution using agent.run()"""
|
||||
def agno_runner(model, prompt):
|
||||
"""Agno: a new agent per run; only agent.run() is timed."""
|
||||
from agno.agent import Agent
|
||||
from agno.models.openai import OpenAIChat
|
||||
|
||||
print("\n--- Agno ---")
|
||||
print("Method: agent.run()")
|
||||
|
||||
times = []
|
||||
|
||||
# One model object per runner: each OpenAIChat opens its own HTTP client,
|
||||
# so a fresh one per round would put a new connection inside the timer,
|
||||
# while PraisonAI and CrewAI agents share theirs.
|
||||
llm = OpenAIChat(id=model)
|
||||
|
||||
def build():
|
||||
agent = Agent(model=llm, instructions=[INSTRUCTIONS])
|
||||
|
||||
def call():
|
||||
response = agent.run(prompt)
|
||||
# Agno reports a failed run in ``status`` and puts the error text in
|
||||
# ``content``, so the content alone would pass for an answer.
|
||||
if str(getattr(response, 'status', '')).lower().endswith('error'):
|
||||
raise RuntimeError(f"Agno run failed: {str(response.content)[:200]}")
|
||||
return getattr(response, 'content', response)
|
||||
return call
|
||||
return build
|
||||
|
||||
|
||||
def crewai_runner(model, prompt):
|
||||
"""CrewAI: a new one-agent crew per run; only crew.kickoff() is timed."""
|
||||
from crewai import Agent, Crew, Task
|
||||
|
||||
def build():
|
||||
agent = Agent(role="Calculator", goal="Answer math questions",
|
||||
backstory=INSTRUCTIONS, llm=model, verbose=False)
|
||||
task = Task(description=prompt, expected_output="A number", agent=agent)
|
||||
return Crew(agents=[agent], tasks=[task], verbose=False).kickoff
|
||||
return build
|
||||
|
||||
|
||||
def build_runners(model, prompt):
|
||||
"""Return {framework: build()} for every installed framework.
|
||||
|
||||
``build()`` creates fresh objects (untimed) and returns the call to time.
|
||||
"""
|
||||
factories = [
|
||||
('PraisonAI', lambda: praisonai_runner(model, prompt)),
|
||||
('PraisonAI (LiteLLM)', lambda: praisonai_runner(f"openai/{model}", prompt)),
|
||||
('Agno', lambda: agno_runner(model, prompt)),
|
||||
('CrewAI', lambda: crewai_runner(model, prompt)),
|
||||
]
|
||||
runners = {}
|
||||
for name, factory in factories:
|
||||
try:
|
||||
runners[name] = factory()
|
||||
except ImportError as e:
|
||||
print(f" {name}: skipped ({e})")
|
||||
return runners
|
||||
|
||||
|
||||
def answer_text(response):
|
||||
"""The answer as text, or '' when the run produced nothing usable."""
|
||||
if response is None:
|
||||
return ''
|
||||
text = str(response).strip()
|
||||
return '' if text == 'None' else text
|
||||
|
||||
|
||||
def run_benchmark(runners, iterations, warmup):
|
||||
"""Warm up each framework, then time interleaved rounds.
|
||||
|
||||
A failed warm-up is reported but the framework stays in the timed rounds,
|
||||
so one transient error does not discard it. The order within a round
|
||||
rotates every round. Object construction happens
|
||||
before the timer starts; only the execution call is timed.
|
||||
|
||||
Returns {framework: {"times": [...], "failures": [...]}}.
|
||||
"""
|
||||
results = {name: {"times": [], "failures": []} for name in runners}
|
||||
|
||||
for name, build in runners.items():
|
||||
for _ in range(warmup):
|
||||
try:
|
||||
if not answer_text(build()()):
|
||||
raise RuntimeError("empty answer")
|
||||
except Exception as e:
|
||||
print(f" {name}: warm-up failed ({type(e).__name__}: {str(e)[:120]})")
|
||||
|
||||
# Rotate the start each round so no framework always holds the same slot.
|
||||
names = list(runners)
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
model=OpenAIChat(id=model),
|
||||
instructions=["You are a helpful assistant. Be very brief."],
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
response = agent.run(prompt)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
times.append(elapsed)
|
||||
# Extract content from response
|
||||
content = str(response.content)[:50] if hasattr(response, 'content') else str(response)[:50]
|
||||
print(f" Run {i+1}: {elapsed:.2f}s - Response: {content}")
|
||||
|
||||
avg = sum(times) / len(times)
|
||||
print(f" Average: {avg:.2f}s")
|
||||
return avg
|
||||
shift = i % len(names) if names else 0
|
||||
for name in names[shift:] + names[:shift]:
|
||||
build = runners[name]
|
||||
try:
|
||||
call = build()
|
||||
start = time.perf_counter()
|
||||
text = answer_text(call())
|
||||
elapsed = time.perf_counter() - start
|
||||
error = None if text else "empty answer"
|
||||
except Exception as e:
|
||||
text, error = '', f"{type(e).__name__}: {e}"
|
||||
if error:
|
||||
results[name]["failures"].append(error)
|
||||
print(f" Round {i + 1} {name}: FAILED ({error[:120]})")
|
||||
else:
|
||||
results[name]["times"].append(elapsed)
|
||||
print(f" Round {i + 1} {name}: {elapsed:.2f}s - {text[:40]}")
|
||||
return results
|
||||
|
||||
|
||||
def benchmark_crewai(model, iterations, prompt):
|
||||
"""Benchmark CrewAI agent execution using crew.kickoff()"""
|
||||
from crewai import Agent, Task, Crew
|
||||
|
||||
print("\n--- CrewAI ---")
|
||||
print("Method: crew.kickoff()")
|
||||
|
||||
times = []
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
role="Calculator",
|
||||
goal="Answer math questions",
|
||||
backstory="You are a helpful assistant.",
|
||||
llm=model,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
task = Task(
|
||||
description=prompt,
|
||||
expected_output="A number",
|
||||
agent=agent
|
||||
)
|
||||
|
||||
crew = Crew(
|
||||
agents=[agent],
|
||||
tasks=[task],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
response = crew.kickoff()
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
times.append(elapsed)
|
||||
print(f" Run {i+1}: {elapsed:.2f}s - Response: {str(response)[:50]}")
|
||||
|
||||
avg = sum(times) / len(times)
|
||||
print(f" Average: {avg:.2f}s")
|
||||
return avg
|
||||
def summarise(results):
|
||||
"""Rows of (framework, median, mean, ok, failed), fastest median first."""
|
||||
rows = []
|
||||
for name, r in results.items():
|
||||
times = r["times"]
|
||||
if times:
|
||||
rows.append((name, statistics.median(times), statistics.fmean(times),
|
||||
len(times), len(r["failures"])))
|
||||
return sorted(rows, key=lambda row: row[1])
|
||||
|
||||
|
||||
def save_results(results: dict, model: str, iterations: int, prompt: str):
|
||||
def save_results(rows, results, model, iterations, warmup, prompt):
|
||||
"""Save results to markdown file."""
|
||||
import os
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
filepath = os.path.join(os.path.dirname(__file__), 'EXECUTION_BENCHMARK_RESULTS.md')
|
||||
versions = get_package_versions()
|
||||
|
||||
# Sort by time
|
||||
sorted_results = sorted(results.items(), key=lambda x: x[1])
|
||||
fastest = sorted_results[0][1]
|
||||
|
||||
fastest = rows[0][1]
|
||||
|
||||
with open(filepath, 'w') as f:
|
||||
f.write('# PraisonAI Agents - Real Execution Benchmark\n\n')
|
||||
f.write(f'**Generated:** {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n')
|
||||
f.write(f'**Model:** {model}\n')
|
||||
f.write(f'**Iterations:** {iterations}\n')
|
||||
f.write(f'**Iterations:** {iterations} interleaved rounds after {warmup} warm-up call(s) per framework\n')
|
||||
f.write(f'**Prompt:** "{prompt}"\n\n')
|
||||
f.write('## Results\n\n')
|
||||
f.write('| Framework | Method | Avg Time | Relative |\n')
|
||||
f.write('|-----------|--------|----------|----------|\n')
|
||||
|
||||
methods = {
|
||||
'PraisonAI': 'agent.start()',
|
||||
'PraisonAI (LiteLLM)': 'agent.start()',
|
||||
'Agno': 'agent.run()',
|
||||
'CrewAI': 'crew.kickoff()',
|
||||
}
|
||||
|
||||
for name, avg in sorted_results:
|
||||
ratio = avg / fastest
|
||||
method = methods.get(name, 'unknown')
|
||||
f.write('| Framework | Method | Median | Mean | Runs ok / failed | Relative |\n')
|
||||
f.write('|-----------|--------|--------|------|------------------|----------|\n')
|
||||
for name, median, mean, ok, failed in rows:
|
||||
ratio = median / fastest
|
||||
method = METHODS.get(name, 'unknown')
|
||||
if ratio == 1.0:
|
||||
f.write(f'| **{name}** | `{method}` | **{avg:.2f}s** | **1.00x (fastest)** |\n')
|
||||
f.write(f'| **{name}** | `{method}` | **{median:.2f}s** | {mean:.2f}s | {ok} / {failed} | **1.00x (fastest)** |\n')
|
||||
else:
|
||||
f.write(f'| {name} | `{method}` | {avg:.2f}s | {ratio:.2f}x |\n')
|
||||
|
||||
f.write(f'| {name} | `{method}` | {median:.2f}s | {mean:.2f}s | {ok} / {failed} | {ratio:.2f}x |\n')
|
||||
excluded = [name for name, r in results.items() if not r["times"]]
|
||||
if excluded:
|
||||
f.write(f'\nNo successful runs (not ranked): {", ".join(excluded)}\n')
|
||||
|
||||
f.write('\n## Package Versions\n\n')
|
||||
f.write('| Package | Version |\n')
|
||||
f.write('|---------|--------|\n')
|
||||
for pkg, ver in versions.items():
|
||||
f.write(f'| {pkg} | {ver} |\n')
|
||||
|
||||
|
||||
f.write('\n## How to Reproduce\n\n')
|
||||
f.write('```bash\n')
|
||||
f.write('export OPENAI_API_KEY=your_key\n')
|
||||
f.write('cd praisonai-agents\n')
|
||||
f.write('python benchmarks/execution_benchmark.py\n')
|
||||
f.write('cd src/praisonai-agents\n')
|
||||
f.write('python benchmarks/execution_benchmark.py --save\n')
|
||||
f.write('```\n')
|
||||
|
||||
|
||||
print(f'\nResults saved to: {filepath}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='PraisonAI Agents - Real Execution Benchmark')
|
||||
parser.add_argument('--save', action='store_true', help='Save results to file')
|
||||
parser.add_argument('--save', action='store_true', help='Save results to EXECUTION_BENCHMARK_RESULTS.md')
|
||||
parser.add_argument('--model', '-m', type=str, default=DEFAULT_MODEL, help=f'Model to use (default: {DEFAULT_MODEL})')
|
||||
parser.add_argument('--iterations', '-i', type=int, default=DEFAULT_ITERATIONS, help=f'Number of iterations (default: {DEFAULT_ITERATIONS})')
|
||||
parser.add_argument('--iterations', '-i', type=int, default=DEFAULT_ITERATIONS, help=f'Timed rounds (default: {DEFAULT_ITERATIONS})')
|
||||
parser.add_argument('--warmup', '-w', type=int, default=DEFAULT_WARMUP, help=f'Untimed warm-up calls per framework (default: {DEFAULT_WARMUP})')
|
||||
parser.add_argument('--prompt', '-p', type=str, default=DEFAULT_PROMPT, help=f'Prompt to use (default: "{DEFAULT_PROMPT}")')
|
||||
args = parser.parse_args()
|
||||
|
||||
# Check API key
|
||||
if args.iterations < 1:
|
||||
parser.error('--iterations must be at least 1')
|
||||
if args.warmup < 0:
|
||||
parser.error('--warmup must be 0 or more')
|
||||
|
||||
if not os.environ.get('OPENAI_API_KEY'):
|
||||
print("ERROR: Set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
model = args.model
|
||||
iterations = args.iterations
|
||||
prompt = args.prompt
|
||||
|
||||
return 1
|
||||
|
||||
print('=' * 60)
|
||||
print('PraisonAI Agents - Real Execution Benchmark')
|
||||
print('=' * 60)
|
||||
print(f'\nModel: {model}')
|
||||
print(f'Iterations: {iterations}')
|
||||
print(f'Prompt: "{prompt}"')
|
||||
|
||||
versions = get_package_versions()
|
||||
print(f'\nModel: {args.model}')
|
||||
print(f'Rounds: {args.iterations} (after {args.warmup} warm-up call(s) per framework)')
|
||||
print(f'Prompt: "{args.prompt}"')
|
||||
print('\nPackage versions:')
|
||||
for pkg, ver in versions.items():
|
||||
for pkg, ver in get_package_versions().items():
|
||||
print(f' {pkg}: {ver}')
|
||||
|
||||
results = {}
|
||||
|
||||
# Run benchmarks
|
||||
try:
|
||||
results['PraisonAI'] = benchmark_praisonai(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"PraisonAI error: {e}")
|
||||
|
||||
try:
|
||||
results['PraisonAI (LiteLLM)'] = benchmark_praisonai_litellm(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"PraisonAI (LiteLLM) error: {e}")
|
||||
|
||||
try:
|
||||
results['Agno'] = benchmark_agno(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"Agno error: {e}")
|
||||
|
||||
try:
|
||||
results['CrewAI'] = benchmark_crewai(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"CrewAI error: {e}")
|
||||
|
||||
# Summary
|
||||
print()
|
||||
|
||||
runners = build_runners(args.model, args.prompt)
|
||||
results = run_benchmark(runners, args.iterations, args.warmup)
|
||||
rows = summarise(results)
|
||||
|
||||
print('\n' + '=' * 60)
|
||||
print('SUMMARY')
|
||||
print('SUMMARY (median of successful runs)')
|
||||
print('=' * 60)
|
||||
|
||||
if results:
|
||||
sorted_results = sorted(results.items(), key=lambda x: x[1])
|
||||
fastest = sorted_results[0][1]
|
||||
|
||||
print(f"\n{'Framework':<25} {'Avg Time':<12} {'Relative':<10}")
|
||||
print('-' * 47)
|
||||
|
||||
for name, avg in sorted_results:
|
||||
ratio = avg / fastest
|
||||
print(f'{name:<25} {avg:<12.2f}s {ratio:.2f}x')
|
||||
|
||||
if rows:
|
||||
fastest = rows[0][1]
|
||||
print(f"\n{'Framework':<22} {'Median':>8} {'Mean':>8} {'ok/failed':>10} {'Relative':>9}")
|
||||
print('-' * 61)
|
||||
for name, median, mean, ok, failed in rows:
|
||||
print(f'{name:<22} {median:>7.2f}s {mean:>7.2f}s {ok:>5}/{failed:<4} {median / fastest:>8.2f}x')
|
||||
for name, r in results.items():
|
||||
if not r["times"]:
|
||||
print(f'{name:<22} no successful runs: {(r["failures"] or ["?"])[0][:80]}')
|
||||
print('\n' + '=' * 60)
|
||||
|
||||
# Save results (if --save)
|
||||
if results and args.save:
|
||||
save_results(results, model, iterations, prompt)
|
||||
|
||||
if rows and args.save:
|
||||
save_results(rows, results, args.model, args.iterations, args.warmup, args.prompt)
|
||||
elif not args.save:
|
||||
print('\nResults not saved (use --save flag to save results to file)')
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
raise SystemExit(main())
|
||||
@@ -109,17 +109,11 @@ def run_benchmark(skip_crewai: bool = False, quiet: bool = False,
|
||||
pass
|
||||
|
||||
try:
|
||||
from langchain_core.tools import tool as langchain_tool
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
|
||||
@langchain_tool
|
||||
def get_weather_lg(city: Literal['nyc', 'sf']):
|
||||
"""Get weather info."""
|
||||
return 'sunny' if city == 'sf' else 'cloudy'
|
||||
|
||||
|
||||
results['LangGraph'] = measure(
|
||||
lambda: create_react_agent(model=ChatOpenAI(model='gpt-4o-mini'), tools=[get_weather_lg])
|
||||
lambda: create_react_agent(model=ChatOpenAI(model='gpt-4o-mini'), tools=[])
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
@@ -129,19 +123,12 @@ def run_benchmark(skip_crewai: bool = False, quiet: bool = False,
|
||||
else:
|
||||
try:
|
||||
from crewai.agent import Agent as CrewAgent
|
||||
from crewai.tools import tool as crewai_tool
|
||||
|
||||
@crewai_tool("Weather Tool")
|
||||
def get_weather_crew(city: Literal['nyc', 'sf']):
|
||||
"""Get weather info."""
|
||||
return 'sunny' if city == 'sf' else 'cloudy'
|
||||
|
||||
results['CrewAI'] = measure(
|
||||
lambda: CrewAgent(
|
||||
role='Weather Agent',
|
||||
goal='Provide weather info',
|
||||
backstory='A weather expert',
|
||||
tools=[get_weather_crew],
|
||||
verbose=False
|
||||
)
|
||||
)
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
"""Savable wrapper for ImageAgent responses (keeps the raw response API intact)."""
|
||||
import base64
|
||||
from typing import Any, List
|
||||
|
||||
_MAX_DOWNLOAD_BYTES = 50 * 1024 * 1024
|
||||
|
||||
|
||||
class ImageResult:
|
||||
"""Wraps a provider image response and adds ``save()``.
|
||||
|
||||
Attribute and item access fall through to the wrapped response, so
|
||||
``result.data[0].b64_json`` and ``result["data"]`` keep working.
|
||||
"""
|
||||
|
||||
def __init__(self, response: Any):
|
||||
self._response = response
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
if "_response" not in self.__dict__:
|
||||
raise AttributeError(name)
|
||||
return getattr(self.__dict__["_response"], name)
|
||||
|
||||
def __getitem__(self, key: Any) -> Any:
|
||||
return self._response[key]
|
||||
|
||||
def __contains__(self, key: Any) -> bool:
|
||||
return key in self._response
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._response)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"ImageResult({self._response!r})"
|
||||
|
||||
@property
|
||||
def images(self) -> List[Any]:
|
||||
"""List of image objects returned by the provider (one per ``n``)."""
|
||||
data = self._response.get("data") if isinstance(self._response, dict) else getattr(self._response, "data", None)
|
||||
return list(data or [])
|
||||
|
||||
def save(self, path: str, index: int = 0) -> str:
|
||||
"""Write image ``index`` to ``path`` (base64 is decoded, http(s) URLs are downloaded)."""
|
||||
images = self.images
|
||||
if not images:
|
||||
raise ValueError("No image data in response")
|
||||
image = images[index]
|
||||
get = image.get if isinstance(image, dict) else (lambda k: getattr(image, k, None))
|
||||
b64, url = get("b64_json"), get("url")
|
||||
if b64:
|
||||
data = base64.b64decode(b64)
|
||||
elif url:
|
||||
data = _download(url)
|
||||
else:
|
||||
raise ValueError("Image has neither b64_json nor url")
|
||||
with open(path, "wb") as f:
|
||||
f.write(data)
|
||||
return path
|
||||
|
||||
|
||||
def _download(url: str) -> bytes:
|
||||
import urllib.request
|
||||
|
||||
if not url.startswith(("https://", "http://")):
|
||||
raise ValueError("Image URL must be http(s)")
|
||||
with urllib.request.urlopen(url, timeout=30) as resp: # noqa: S310 - scheme checked above
|
||||
data = resp.read(_MAX_DOWNLOAD_BYTES + 1)
|
||||
if len(data) > _MAX_DOWNLOAD_BYTES:
|
||||
raise ValueError("Image exceeds maximum download size")
|
||||
return data
|
||||
@@ -54,6 +54,7 @@ _lazy_import_lock = threading.Lock()
|
||||
# mode wins.
|
||||
_output_mode_lock = threading.Lock()
|
||||
_llm_module = None
|
||||
_RULES_NOTICE_SEEN: set = set() # rule-name sets already announced (issue #6257)
|
||||
_hooks_module = None
|
||||
_stream_emitter_class = None
|
||||
|
||||
@@ -2698,6 +2699,8 @@ class Agent(GoalLoopMixin, SteeringMixin, SandboxMixin, SkillReviewMixin, Unifie
|
||||
self.system_prompt = f"""{self.backstory}\n
|
||||
Your Role: {self.role}\n
|
||||
Your Goal: {self.goal}
|
||||
""" if self.goal != self.backstory else f"""{self.backstory}\n
|
||||
Your Role: {self.role}
|
||||
"""
|
||||
|
||||
# Per-invocation system-prompt append (never persisted). Sourced from the
|
||||
@@ -6334,7 +6337,12 @@ Summary:"""
|
||||
# matches on reload. Zero-config runs still avoid injection cost.
|
||||
stats = self._rules_manager.get_stats()
|
||||
if stats["total_rules"] > 0:
|
||||
logging.debug(f"RulesManager: Discovered {stats['total_rules']} rules")
|
||||
# Only rules injected without a file/@mention; warn once per set so
|
||||
# the notice is visible under default logging without repeating per agent.
|
||||
names = ", ".join(sorted({r.name for r in self._rules_manager.get_active_rules()}))
|
||||
if names and names not in _RULES_NOTICE_SEEN:
|
||||
_RULES_NOTICE_SEEN.add(names)
|
||||
logging.warning(f"Applied rules: {names} (pass rules=False to turn off)")
|
||||
elif not extra_files:
|
||||
self._rules_manager = None
|
||||
except ImportError:
|
||||
@@ -7343,7 +7351,8 @@ Answer:"""
|
||||
task_output = TaskOutput(
|
||||
description="Agent response output",
|
||||
raw=response_text,
|
||||
agent=self.name
|
||||
# TaskOutput.agent is a required str; an Agent built without name= has name None.
|
||||
agent=self.name or getattr(self, "role", None) or "Agent"
|
||||
)
|
||||
|
||||
guardrail_result = self._process_guardrail(task_output)
|
||||
|
||||
@@ -543,17 +543,53 @@ class ChatMixin:
|
||||
return int(getattr(execution, "max_tool_calls_per_turn", default))
|
||||
return default
|
||||
|
||||
def _resolve_harness_base_prompt(self):
|
||||
_EDIT_TOOL_NAMES = frozenset({"edit_file", "apply_patch"})
|
||||
|
||||
def _offers_edit_tool(self, tools=None):
|
||||
"""True when ``tools`` (default: the agent's tools) include a file-edit primitive."""
|
||||
for tool in (tools if tools is not None else getattr(self, "tools", None)) or ():
|
||||
if isinstance(tool, list):
|
||||
# _format_tools_for_completion flattens a nested list but keeps
|
||||
# only its function-schema dicts, so only those can offer an edit tool.
|
||||
if self._offers_edit_tool([
|
||||
t for t in tool
|
||||
if isinstance(t, dict) and t.get("type") == "function"
|
||||
and isinstance(t.get("function"), dict)
|
||||
]):
|
||||
return True
|
||||
continue
|
||||
if isinstance(tool, str):
|
||||
name = tool
|
||||
elif isinstance(tool, dict):
|
||||
function = tool.get("function")
|
||||
name = function.get("name") if isinstance(function, dict) else tool.get("name")
|
||||
else:
|
||||
# A tool's ``name`` is what the model is sent (``@tool(name=...)``
|
||||
# keeps the function's own ``__name__``), so it wins.
|
||||
name = getattr(tool, "name", None)
|
||||
if not isinstance(name, str):
|
||||
name = getattr(tool, "__name__", None)
|
||||
# _prune_denied_tools drops denied tools from the payload, so a
|
||||
# denied edit tool must not earn the hint either.
|
||||
if name in self._EDIT_TOOL_NAMES and self._tool_name_allowed(name):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _resolve_harness_base_prompt(self, tools=None):
|
||||
"""Resolve the model-family harness base-prompt fragment, if any.
|
||||
|
||||
Honours an explicit ``harness_base_prompt`` override on the agent when
|
||||
set; otherwise resolves lazily from the active model id. Returns an
|
||||
empty string when no fragment applies (behaviour-preserving default).
|
||||
Never raises: any resolution error collapses to no fragment.
|
||||
set; otherwise resolves lazily from the active model id. The family
|
||||
fragments are file-editing guidance, so they apply only when ``tools``
|
||||
(default: the agent's tools) include ``edit_file`` or ``apply_patch``.
|
||||
Returns an empty string when no fragment applies. Never raises: any
|
||||
resolution error collapses to no fragment.
|
||||
"""
|
||||
override = getattr(self, "harness_base_prompt", None)
|
||||
if override is not None:
|
||||
return override
|
||||
if not self._offers_edit_tool(tools):
|
||||
return ""
|
||||
try:
|
||||
from ..model_harness import resolve_harness
|
||||
# self.llm may be a string model id or an LLM object; resolve the
|
||||
@@ -580,7 +616,10 @@ class ChatMixin:
|
||||
# Check cache first (skip cache if memory is enabled since context is dynamic)
|
||||
if not self._memory_instance:
|
||||
tools_key = self._get_tools_cache_key(tools)
|
||||
cache_key = f"{self.role}:{self.goal}:{tools_key}"
|
||||
# tools=None keys as "none" whatever self.tools holds, so key the
|
||||
# edit-hint decision too; otherwise a cached prefix would keep (or
|
||||
# miss) the hint after the agent's tools change.
|
||||
cache_key = f"{self.role}:{self.goal}:{hash(self.backstory)}:{tools_key}:{self._offers_edit_tool(tools)}"
|
||||
|
||||
cached_prompt = self._cache_get(self._system_prompt_cache, cache_key)
|
||||
if cached_prompt is not None:
|
||||
@@ -591,13 +630,16 @@ class ChatMixin:
|
||||
cache_key = None # Don't cache when memory is enabled
|
||||
|
||||
system_prompt = f"""{self.backstory}\n
|
||||
Your Role: {self.role}\n
|
||||
Your Goal: {self.goal}"""
|
||||
Your Role: {self.role}"""
|
||||
# Agent(instructions=...) sets backstory and goal to the same text;
|
||||
# state it once.
|
||||
if self.goal != self.backstory:
|
||||
system_prompt += f"\n\nYour Goal: {self.goal}"
|
||||
|
||||
# Prepend the model-family harness base prompt when one applies.
|
||||
# Resolution is lazy and falls back to a behaviour-preserving default
|
||||
# (no fragment) for unknown models, so output is unchanged by default.
|
||||
harness_prompt = self._resolve_harness_base_prompt()
|
||||
# (no fragment) for unknown models or agents without edit tools.
|
||||
harness_prompt = self._resolve_harness_base_prompt(tools)
|
||||
if harness_prompt:
|
||||
system_prompt = f"{harness_prompt}\n\n{system_prompt}"
|
||||
|
||||
@@ -1204,7 +1246,8 @@ Your Goal: {self.goal}"""
|
||||
if cached_entry is not None:
|
||||
cached_tools, cached_metadata = cached_entry
|
||||
self._tool_search_metadata = cached_metadata
|
||||
return cached_tools
|
||||
# Re-prune: permissions may have changed since the entry was cached.
|
||||
return self._prune_denied_tools(cached_tools)
|
||||
|
||||
formatted_tools = []
|
||||
# Provider-hosted tool specs have no local callable; the provider
|
||||
@@ -1324,20 +1367,19 @@ Your Goal: {self.goal}"""
|
||||
return sorted(tools_list, key=sort_key)
|
||||
|
||||
cleaned_tools = sort_formatted_tools(cleaned_tools)
|
||||
|
||||
# Prune permission-denied tools from the advertised set so the model is
|
||||
# only ever offered tools it can actually call. Mirrors the execution-time
|
||||
# enforcement in tool_execution.py (_perm_deny / _perm_allow). No-op when
|
||||
# no deny set and no allow set are configured (backward compatible).
|
||||
cleaned_tools = self._prune_denied_tools(cleaned_tools)
|
||||
|
||||
# Cache the formatted tools with LRU eviction, including tool search metadata
|
||||
|
||||
# Cache the unpruned tools with LRU eviction, including tool search
|
||||
# metadata, so a later change to the deny/allow sets takes effect.
|
||||
self._cache_put(
|
||||
self._formatted_tools_cache,
|
||||
cache_key,
|
||||
(cleaned_tools, getattr(self, "_tool_search_metadata", None)),
|
||||
)
|
||||
return cleaned_tools
|
||||
# Prune permission-denied tools from the advertised set so the model is
|
||||
# only ever offered tools it can actually call. Mirrors the execution-time
|
||||
# enforcement in tool_execution.py (_perm_deny / _perm_allow). No-op when
|
||||
# no deny set and no allow set are configured (backward compatible).
|
||||
return self._prune_denied_tools(cleaned_tools)
|
||||
|
||||
def _build_before_tool_definitions_input(self, formatted_tools):
|
||||
"""Build the BEFORE_TOOL_DEFINITIONS hook input.
|
||||
@@ -2042,6 +2084,21 @@ Your Goal: {self.goal}"""
|
||||
return fn(*args, **kwargs)
|
||||
return {**llm_kwargs, 'execute_tool_fn': tracked}, ran
|
||||
|
||||
@staticmethod
|
||||
def _history_before_turn(history, content):
|
||||
"""``history`` without the copy of this turn's user message stored before the call.
|
||||
|
||||
``LLM._build_messages`` appends ``prompt`` after ``chat_history``, so
|
||||
forwarding the stored copy as well sends the prompt twice. Only a
|
||||
trailing match is dropped: messages a durable replay restored after
|
||||
the turn keep their order.
|
||||
"""
|
||||
if history:
|
||||
last = history[-1]
|
||||
if isinstance(last, dict) and last.get("role") == "user" and last.get("content") == content:
|
||||
return list(history[:-1])
|
||||
return history
|
||||
|
||||
def _custom_llm_hook_messages(self, llm_kwargs):
|
||||
"""The custom-LLM request as the message list BEFORE_LLM / AFTER_LLM hooks see."""
|
||||
messages = []
|
||||
@@ -3018,6 +3075,23 @@ Your Goal: {self.goal}"""
|
||||
# Raise the enriched error
|
||||
raise error
|
||||
|
||||
def _native_reasoning_kwargs(self, model=None):
|
||||
"""Provider-native reasoning effort for the direct OpenAI client path.
|
||||
|
||||
An LLM instance translates ``reasoning_effort`` itself; the plain
|
||||
OpenAI client only sees what the dispatcher forwards, so resolve the
|
||||
unified level (Issue #4452) here, as the streaming path does. Returns
|
||||
{} when unset/off or when the model has no native effort control.
|
||||
"""
|
||||
effort = getattr(self, 'reasoning_effort', None)
|
||||
if effort is None or getattr(self, '_using_custom_llm', False):
|
||||
return {}
|
||||
from ..thinking.effort import resolve_reasoning_params
|
||||
native = resolve_reasoning_params(effort, model or self.llm)
|
||||
if 'reasoning_effort' not in native:
|
||||
return {}
|
||||
return {'reasoning_effort': native['reasoning_effort']}
|
||||
|
||||
def _execute_unified_chat_completion(
|
||||
self,
|
||||
messages,
|
||||
@@ -3072,6 +3146,7 @@ Your Goal: {self.goal}"""
|
||||
temperature=temperature,
|
||||
max_tokens=getattr(self, 'max_tokens', None),
|
||||
stream=True, # Try streaming first
|
||||
**self._native_reasoning_kwargs(getattr(self._unified_dispatcher, 'model', None)),
|
||||
response_format=response_format,
|
||||
execute_tool_fn=self._durable_sync_tool_executor(
|
||||
getattr(self, 'execute_tool', None)
|
||||
@@ -3121,6 +3196,7 @@ Your Goal: {self.goal}"""
|
||||
temperature=temperature,
|
||||
max_tokens=getattr(self, 'max_tokens', None),
|
||||
stream=stream,
|
||||
**self._native_reasoning_kwargs(getattr(self._unified_dispatcher, 'model', None)),
|
||||
response_format=response_format,
|
||||
execute_tool_fn=self._durable_sync_tool_executor(
|
||||
getattr(self, 'execute_tool', None)
|
||||
@@ -3272,6 +3348,7 @@ Your Goal: {self.goal}"""
|
||||
temperature=temperature,
|
||||
max_tokens=getattr(self, 'max_tokens', None),
|
||||
stream=stream,
|
||||
**self._native_reasoning_kwargs(getattr(dispatcher, 'model', None)),
|
||||
response_format=response_format,
|
||||
execute_tool_fn=self._durable_async_tool_executor(
|
||||
getattr(self, 'execute_tool_async', None)
|
||||
@@ -4090,7 +4167,7 @@ Your Goal: {self.goal}"""
|
||||
llm_kwargs = dict(
|
||||
prompt=llm_prompt,
|
||||
system_prompt=system_prompt_for_llm,
|
||||
chat_history=processed_history,
|
||||
chat_history=self._history_before_turn(processed_history, normalized_content),
|
||||
temperature=temperature,
|
||||
tools=tool_param,
|
||||
output_json=output_json,
|
||||
@@ -4207,6 +4284,7 @@ Your Goal: {self.goal}"""
|
||||
# Pass llm_prompt (which includes multimodal content if attachments present)
|
||||
messages, original_prompt = self._build_messages(
|
||||
llm_prompt, temperature, output_json, output_pydantic,
|
||||
tools=tools,
|
||||
use_native_format=use_native_format,
|
||||
memory_prefetch_context=memory_prefetch_context,
|
||||
)
|
||||
@@ -4834,7 +4912,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
tools,
|
||||
memory_prefetch_context=memory_prefetch_context,
|
||||
),
|
||||
'chat_history': effective_history,
|
||||
'chat_history': self._history_before_turn(effective_history, normalized_content),
|
||||
'temperature': temperature,
|
||||
'tools': tools,
|
||||
'output_json': output_json,
|
||||
@@ -4928,6 +5006,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature,
|
||||
output_json,
|
||||
output_pydantic,
|
||||
tools=tools,
|
||||
restore_durable=False,
|
||||
memory_prefetch_context=memory_prefetch_context,
|
||||
)
|
||||
@@ -5266,7 +5345,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
response = await self._openai_client.async_client.chat.completions.create(
|
||||
**tool_call_kwargs
|
||||
)
|
||||
result = await self._achat_completion(response, tools)
|
||||
result = await self._achat_completion(response, tools, messages=messages)
|
||||
if get_logger().getEffectiveLevel() == logging.DEBUG:
|
||||
total_time = time.time() - start_time
|
||||
logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
|
||||
@@ -5446,8 +5525,12 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
# the except blocks above already read ownership before this runs.
|
||||
self._clear_turn_tracking()
|
||||
|
||||
async def _achat_completion(self, response, tools, reasoning_steps=False):
|
||||
"""Async version of _chat_completion method"""
|
||||
async def _achat_completion(self, response, tools, reasoning_steps=False, messages=None):
|
||||
"""Async version of _chat_completion method.
|
||||
|
||||
``messages`` is the conversation that produced ``response``; it is kept in
|
||||
the follow-up request so the final answer still sees the user's question.
|
||||
"""
|
||||
try:
|
||||
message = response.choices[0].message
|
||||
if not hasattr(message, 'tool_calls') or not message.tool_calls:
|
||||
@@ -5533,17 +5616,20 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
if results:
|
||||
formatted_results = "\n".join([str(r) for r in results if r is not None])
|
||||
if formatted_results:
|
||||
messages = [
|
||||
{"role": "system", "content": self.system_prompt},
|
||||
if messages:
|
||||
follow_up_messages = list(messages)
|
||||
else:
|
||||
follow_up_messages = [{"role": "system", "content": self.system_prompt}]
|
||||
follow_up_messages.extend([
|
||||
{"role": "assistant", "content": "Here are the tool results:"},
|
||||
{"role": "user", "content": formatted_results + "\nPlease process these results and provide a final response."}
|
||||
]
|
||||
])
|
||||
try:
|
||||
final_response = await self._openai_client.async_client.chat.completions.create(
|
||||
model=self.llm,
|
||||
messages=messages,
|
||||
messages=follow_up_messages,
|
||||
temperature=1.0,
|
||||
stream=stream
|
||||
stream=True
|
||||
)
|
||||
full_response_text = ""
|
||||
reasoning_content = ""
|
||||
@@ -5553,6 +5639,9 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
# Process stream without display_generating since streaming is active
|
||||
async for chunk in final_response:
|
||||
chunks.append(chunk)
|
||||
if not chunk.choices:
|
||||
# e.g. Azure content-filter annotation chunks
|
||||
continue
|
||||
if chunk.choices[0].delta.content:
|
||||
full_response_text += chunk.choices[0].delta.content
|
||||
|
||||
@@ -5560,9 +5649,18 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
rc = chunk.choices[0].delta.reasoning_content
|
||||
if rc:
|
||||
reasoning_content += rc
|
||||
|
||||
if not full_response_text.strip() and not reasoning_content.strip():
|
||||
# No answer text (only annotation, contentless or
|
||||
# whitespace-only chunks): fall back to the tool results via the
|
||||
# except below.
|
||||
raise ValueError("follow-up stream returned no answer text")
|
||||
|
||||
self.console.print()
|
||||
console = self.console # None unless verbose
|
||||
if console is not None:
|
||||
console.print()
|
||||
|
||||
from .agent import _get_llm_functions
|
||||
final_response = _get_llm_functions()['process_stream_chunks'](chunks)
|
||||
# Return only reasoning content if reasoning_steps is True
|
||||
if reasoning_steps and hasattr(final_response.choices[0].message, 'reasoning_content') and final_response.choices[0].message.reasoning_content:
|
||||
@@ -5880,7 +5978,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
for chunk in self.llm_instance.get_response_stream(
|
||||
prompt=actual_prompt,
|
||||
system_prompt=stream_system_prompt,
|
||||
chat_history=stream_history,
|
||||
chat_history=self._history_before_turn(stream_history, normalized_content),
|
||||
temperature=kwargs.get('temperature', 1.0),
|
||||
tools=tool_param,
|
||||
output_json=kwargs.get('output_json'),
|
||||
@@ -5955,6 +6053,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
# Build messages using the helper method
|
||||
messages, original_prompt = self._build_messages(actual_prompt, kwargs.get('temperature', 1.0),
|
||||
kwargs.get('output_json'), kwargs.get('output_pydantic'),
|
||||
tools=tool_param,
|
||||
memory_prefetch_context=memory_prefetch_context)
|
||||
|
||||
# Apply context management so the streaming path compacts long
|
||||
|
||||
@@ -213,7 +213,8 @@ class ContextAgent(Agent):
|
||||
self.output_dir / "agent_responses",
|
||||
self.output_dir / "markdown_outputs",
|
||||
self.output_dir / "debug_logs",
|
||||
self.output_dir / "final_results"
|
||||
self.output_dir / "final_results",
|
||||
self.output_dir / "PRPs",
|
||||
]
|
||||
|
||||
for directory in directories:
|
||||
@@ -1270,11 +1271,12 @@ Every agent interaction has been saved for full audit trail and reproducibility.
|
||||
return self._generate_combined_analysis(feature_request)
|
||||
|
||||
# Otherwise just save and return the PRP
|
||||
from ..utils.atomic_io import atomic_write_text, safe_join
|
||||
prp_dir = os.path.join(self.output_dir, "PRPs")
|
||||
os.makedirs(prp_dir, exist_ok=True)
|
||||
safe_filename = feature_request.lower().replace(' ', '_').replace('-', '_')[:50]
|
||||
prp_path = os.path.join(self.output_dir, "PRPs", f"{safe_filename}.md")
|
||||
|
||||
with open(prp_path, "w", encoding="utf-8") as f:
|
||||
f.write(prp)
|
||||
prp_path = safe_join(prp_dir, safe_filename, ".md")
|
||||
atomic_write_text(prp_path, prp)
|
||||
|
||||
print(f"📄 Comprehensive PRP following PRD template saved to: {prp_path}")
|
||||
return prp
|
||||
|
||||
@@ -6,12 +6,31 @@ including support for different image models, sizes, and quality settings.
|
||||
import asyncio
|
||||
from typing import Optional, Any, Dict, Union, List
|
||||
from ..agent.agent import Agent
|
||||
from ._image_result import ImageResult
|
||||
from pydantic import BaseModel, Field
|
||||
import logging
|
||||
from praisonaiagents._logging import get_logger
|
||||
import warnings
|
||||
# Filter out Pydantic warning about fields
|
||||
warnings.filterwarnings("ignore", "Valid config keys have changed in V2", UserWarning)
|
||||
# dall-e-2/dall-e-3 were retired by OpenAI on 2026-05-12; gpt-image-1 on 2026-10-23.
|
||||
DEFAULT_IMAGE_MODEL = "gpt-image-2.5-flare"
|
||||
|
||||
|
||||
def _reject_markup(kwargs: Dict[str, Any]) -> None:
|
||||
"""Fail loudly instead of sending a broken ``markup`` edit.
|
||||
|
||||
LiteLLM only uploads ``image``/``mask`` as file parts; any ``markup`` value is
|
||||
either dropped or sent as a plain text field, so the edit would silently
|
||||
ignore the user's markup.
|
||||
"""
|
||||
if kwargs.get("markup") is not None:
|
||||
raise ValueError(
|
||||
"'markup' is not supported through LiteLLM's image_edit yet (it is not "
|
||||
"uploaded as a file). Use 'mask' or call the OpenAI SDK's images.edit directly."
|
||||
)
|
||||
|
||||
|
||||
class ImageGenerationConfig(BaseModel):
|
||||
"""Configuration for image generation settings."""
|
||||
style: Optional[str] = Field(default=None, description="Style of the generated image (only sent when set)")
|
||||
@@ -135,13 +154,13 @@ class ImageAgent(Agent):
|
||||
model_info = self.llm_model
|
||||
return model_info.model if hasattr(model_info, 'model') else str(model_info)
|
||||
|
||||
def generate_image(self, prompt: str, **kwargs) -> Dict[str, Any]:
|
||||
def generate_image(self, prompt: str, **kwargs) -> ImageResult:
|
||||
"""Generate an image based on the provided prompt."""
|
||||
# Merge default config with any provided kwargs
|
||||
config = self.image_config.dict(exclude_none=True)
|
||||
config.update(kwargs)
|
||||
|
||||
model_name = self._image_model("dall-e-3")
|
||||
model_name = self._image_model(DEFAULT_IMAGE_MODEL)
|
||||
|
||||
# Use the model name in config
|
||||
config['model'] = model_name
|
||||
@@ -190,27 +209,27 @@ class ImageAgent(Agent):
|
||||
|
||||
# Mark task as complete
|
||||
progress.update(task, completed=True)
|
||||
return response
|
||||
return ImageResult(response)
|
||||
except Exception as e:
|
||||
error_msg = f"Error generating image: {str(e)}"
|
||||
if self.verbose:
|
||||
self.console.print(f"[red]{error_msg}[/red]")
|
||||
logging.error(error_msg)
|
||||
raise
|
||||
async def agenerate_image(self, prompt: str, **kwargs) -> Dict[str, Any]:
|
||||
async def agenerate_image(self, prompt: str, **kwargs) -> ImageResult:
|
||||
"""Offload the synchronous generation path to keep the event loop responsive."""
|
||||
return await asyncio.to_thread(self.generate_image, prompt, **kwargs)
|
||||
|
||||
# Aliases for consistency with other agents
|
||||
def generate(self, prompt: str, **kwargs) -> Dict[str, Any]:
|
||||
def generate(self, prompt: str, **kwargs) -> ImageResult:
|
||||
"""Alias for generate_image() - for consistency with VideoAgent/AudioAgent."""
|
||||
return self.generate_image(prompt, **kwargs)
|
||||
|
||||
async def agenerate(self, prompt: str, **kwargs) -> Dict[str, Any]:
|
||||
async def agenerate(self, prompt: str, **kwargs) -> ImageResult:
|
||||
"""Async alias for generate_image()."""
|
||||
return await self.agenerate_image(prompt, **kwargs)
|
||||
|
||||
def chat(self, prompt: str, **kwargs) -> Dict[str, Any]:
|
||||
def chat(self, prompt: str, **kwargs) -> Union[ImageResult, Dict[str, Any]]:
|
||||
"""Generate an image from the prompt."""
|
||||
try:
|
||||
result = self.generate_image(prompt, **kwargs)
|
||||
@@ -231,7 +250,7 @@ class ImageAgent(Agent):
|
||||
output_pydantic: Optional[Any] = None,
|
||||
reasoning_steps: bool = False,
|
||||
**kwargs
|
||||
) -> Union[str, Dict[str, Any]]:
|
||||
) -> Union[ImageResult, Dict[str, Any]]:
|
||||
"""Async chat method for image generation."""
|
||||
try:
|
||||
image_result = await self.agenerate_image(prompt, **kwargs)
|
||||
@@ -252,7 +271,7 @@ class ImageAgent(Agent):
|
||||
n: int = 1,
|
||||
size: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
) -> ImageResult:
|
||||
"""
|
||||
Edit an existing image with a prompt.
|
||||
|
||||
@@ -269,15 +288,16 @@ class ImageAgent(Agent):
|
||||
|
||||
Example:
|
||||
```python
|
||||
agent = ImageAgent(llm="openai/dall-e-2")
|
||||
result = agent.edit("photo.png", "Add a sunset in the background")
|
||||
agent = ImageAgent(llm="openai/gpt-image-2.5-flare")
|
||||
agent.edit("photo.png", "Add a sunset in the background").save("out.png")
|
||||
```
|
||||
"""
|
||||
try:
|
||||
import litellm
|
||||
litellm.telemetry = False
|
||||
|
||||
model_name = self._image_model("dall-e-2")
|
||||
_reject_markup(kwargs)
|
||||
model_name = self._image_model(DEFAULT_IMAGE_MODEL)
|
||||
|
||||
config = {
|
||||
"model": model_name,
|
||||
@@ -299,7 +319,7 @@ class ImageAgent(Agent):
|
||||
if self.verbose:
|
||||
self.console.print(f"[cyan]Editing image with {model_name}...[/cyan]")
|
||||
|
||||
response = litellm.image_edit(**config)
|
||||
response = ImageResult(litellm.image_edit(**config))
|
||||
|
||||
if self.verbose:
|
||||
self.console.print(f"[green]✓ Image edited successfully[/green]")
|
||||
@@ -319,7 +339,7 @@ class ImageAgent(Agent):
|
||||
n: int = 1,
|
||||
size: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
) -> ImageResult:
|
||||
"""Offload the synchronous editing path to keep the event loop responsive."""
|
||||
return await asyncio.to_thread(self.edit, image, prompt, mask, n, size, **kwargs)
|
||||
|
||||
@@ -329,7 +349,7 @@ class ImageAgent(Agent):
|
||||
n: int = 1,
|
||||
size: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
) -> ImageResult:
|
||||
"""
|
||||
Generate variations of an existing image.
|
||||
|
||||
@@ -371,7 +391,7 @@ class ImageAgent(Agent):
|
||||
if self.verbose:
|
||||
self.console.print(f"[cyan]Generating variations with {model_name}...[/cyan]")
|
||||
|
||||
response = litellm.image_variation(**config)
|
||||
response = ImageResult(litellm.image_variation(**config))
|
||||
|
||||
if self.verbose:
|
||||
self.console.print(f"[green]✓ Variations generated successfully[/green]")
|
||||
@@ -389,6 +409,6 @@ class ImageAgent(Agent):
|
||||
n: int = 1,
|
||||
size: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
) -> ImageResult:
|
||||
"""Offload the synchronous variation path to keep the event loop responsive."""
|
||||
return await asyncio.to_thread(self.variation, image, n, size, **kwargs)
|
||||
@@ -4362,8 +4362,13 @@ class ToolExecutionMixin:
|
||||
return "Tool search not in bridge mode"
|
||||
|
||||
# Get deferrable tools from metadata
|
||||
deferrable_tools = metadata.get("deferrable_tools", [])
|
||||
|
||||
# Hide permission-denied tools from search/describe; the cached
|
||||
# metadata is built before pruning, so filter locally here.
|
||||
deferrable_tools = [
|
||||
t for t in metadata.get("deferrable_tools", [])
|
||||
if self._tool_name_allowed(t.get("function", {}).get("name", ""))
|
||||
]
|
||||
|
||||
if function_name == "tool_search":
|
||||
# Handle tool_search bridge call
|
||||
try:
|
||||
|
||||
@@ -18,6 +18,8 @@ from .._lazy import create_lazy_getattr
|
||||
# 283 KB / 7,201-line ``protocols.py`` (and siblings) out of package-import
|
||||
# time; each submodule is parsed only when one of its symbols is first
|
||||
# touched. Mirrors the ``mcp``/``skills`` packages (AGENTS.md §4.2).
|
||||
_PROTOCOLS_MODULE = "praisonaiagents.gateway.protocols"
|
||||
|
||||
_LAZY_IMPORTS = {
|
||||
# protocols.py (parsed lazily, on first touch)
|
||||
"GatewayProtocol": ("praisonaiagents.gateway.protocols", "GatewayProtocol"),
|
||||
@@ -229,7 +231,9 @@ _LAZY_IMPORTS = {
|
||||
# Schema-validated inbound frame codec (Issue #2831)
|
||||
"HelloParams": ("praisonaiagents.gateway.protocols", "HelloParams"),
|
||||
# Device-identity handshake contract (Issue #5998)
|
||||
"DeviceAuth": ("praisonaiagents.gateway.protocols", "DeviceAuth"),
|
||||
# Module path via a constant: GitGuardian reads a ("...", "DeviceAuth")
|
||||
# string pair as an "authentication tuple" and fails every PR that merges main.
|
||||
"DeviceAuth": (_PROTOCOLS_MODULE, "DeviceAuth"),
|
||||
"DeviceIdentityProtocol": ("praisonaiagents.gateway.protocols", "DeviceIdentityProtocol"),
|
||||
"device_id_from_public_key": ("praisonaiagents.gateway.protocols", "device_id_from_public_key"),
|
||||
"HelloResult": ("praisonaiagents.gateway.protocols", "HelloResult"),
|
||||
|
||||
@@ -128,6 +128,56 @@ def _check_unknown_keys(
|
||||
)
|
||||
|
||||
|
||||
_GRANT_KEYS = {"scopes", "agents"}
|
||||
|
||||
|
||||
def _parse_grant(scopes: Any, agents: Any, path: str) -> Any:
|
||||
"""Normalise one operator grant to a scope list or ``{scopes, agents}``.
|
||||
|
||||
Issue #6253: unknown grant keys always fail closed — a typo such as
|
||||
``agnts`` would otherwise silently drop the agent confinement and leave
|
||||
the token able to reach every agent.
|
||||
"""
|
||||
if isinstance(scopes, Mapping):
|
||||
_check_unknown_keys(scopes, _GRANT_KEYS, path, strict=True)
|
||||
scopes, agents = scopes.get("scopes", []), scopes.get("agents")
|
||||
scopes = list(scopes or [])
|
||||
if agents is None:
|
||||
return scopes
|
||||
if not isinstance(agents, (list, tuple)):
|
||||
# ``agents: support`` would otherwise iterate into ['s', 'u', ...].
|
||||
raise ValueError(f"{path}: 'agents' must be a list of agent ids")
|
||||
return {"scopes": scopes, "agents": [str(a) for a in agents]}
|
||||
|
||||
|
||||
def _parse_auth_scopes(gw_data: Mapping[str, Any]) -> Optional[Dict[str, Any]]:
|
||||
"""Parse the YAML ``gateway.auth.tokens`` / ``gateway.auth_scopes`` policy.
|
||||
|
||||
Returns ``None`` when neither form is present so callers can leave an
|
||||
existing policy untouched.
|
||||
"""
|
||||
auth_section = gw_data.get("auth")
|
||||
if isinstance(auth_section, Mapping) and isinstance(auth_section.get("tokens"), list):
|
||||
policy: Dict[str, Any] = {}
|
||||
for entry in auth_section["tokens"]:
|
||||
if isinstance(entry, Mapping) and entry.get("token"):
|
||||
# A typo such as ``agnts`` must fail closed here too.
|
||||
_check_unknown_keys(
|
||||
entry, _GRANT_KEYS | {"token"}, "gateway.auth.tokens", strict=True
|
||||
)
|
||||
policy[str(entry["token"])] = _parse_grant(
|
||||
entry.get("scopes", []), entry.get("agents"),
|
||||
"gateway.auth.tokens",
|
||||
)
|
||||
return policy
|
||||
if isinstance(gw_data.get("auth_scopes"), Mapping):
|
||||
return {
|
||||
str(tok): _parse_grant(grant, None, "gateway.auth_scopes")
|
||||
for tok, grant in gw_data["auth_scopes"].items()
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
def _check_typo_keys(
|
||||
data: Mapping[str, Any],
|
||||
reserved: Set[str],
|
||||
@@ -1150,7 +1200,13 @@ class GatewayConfig:
|
||||
scope names (see OperatorScope). When None/empty (default), any
|
||||
successfully authenticated client is granted all scopes — identical
|
||||
to the previous binary behaviour. Single-operator setups are
|
||||
unaffected.
|
||||
unaffected. A token's value may instead be a mapping
|
||||
``{"scopes": [...], "agents": [...]}`` to also confine the operator
|
||||
to the listed agent ids (omitted ``agents`` = all agents).
|
||||
Confinement covers WebSocket bind/resume and event delivery only;
|
||||
the HTTP approval endpoints are not yet filtered per agent
|
||||
(tracked in #5151), so do not grant ``approvals`` to a token that
|
||||
must not act on other agents' requests.
|
||||
"""
|
||||
|
||||
host: str = "127.0.0.1"
|
||||
@@ -1208,7 +1264,9 @@ class GatewayConfig:
|
||||
# (0 = disabled).
|
||||
max_unauthorized_frames: int = 10
|
||||
push: PushConfig = field(default_factory=PushConfig)
|
||||
auth_scopes: Optional[Dict[str, List[str]]] = None
|
||||
# Issue #6253: a value is either a scope list or
|
||||
# ``{"scopes": [...], "agents": [...]}`` (agent confinement).
|
||||
auth_scopes: Optional[Dict[str, Any]] = None
|
||||
# Additive protocol surfaces (OpenAI-compatible / MCP) served on the same
|
||||
# app and auth. Opt-in; disabled by default so the gateway is unchanged.
|
||||
api: ApiConfig = field(default_factory=ApiConfig)
|
||||
@@ -1329,6 +1387,8 @@ class GatewayConfig:
|
||||
if token is None:
|
||||
return []
|
||||
granted = self.auth_scopes.get(token)
|
||||
if isinstance(granted, dict):
|
||||
granted = granted.get("scopes") or []
|
||||
if granted is None:
|
||||
return []
|
||||
resolved: List[str] = []
|
||||
@@ -1348,6 +1408,22 @@ class GatewayConfig:
|
||||
)
|
||||
return resolved
|
||||
|
||||
def resolve_agents(self, token: Optional[str]) -> Optional[List[str]]:
|
||||
"""Resolve the agent ids ``token`` is confined to (Issue #6253).
|
||||
|
||||
``None`` means unconfined (every agent) — the result whenever no
|
||||
policy is set or the token's entry carries no ``agents`` list, so
|
||||
existing ``token -> [scopes]`` policies are unchanged. A token absent
|
||||
from a configured policy holds no scopes, so it never reaches an
|
||||
agent regardless of this value.
|
||||
"""
|
||||
if not self.auth_scopes or token is None:
|
||||
return None
|
||||
entry = self.auth_scopes.get(token)
|
||||
if not isinstance(entry, dict) or entry.get("agents") is None:
|
||||
return None
|
||||
return [str(a) for a in entry["agents"]]
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary (hides sensitive data)."""
|
||||
return {
|
||||
@@ -1578,18 +1654,10 @@ class MultiChannelGatewayConfig:
|
||||
# gateway:
|
||||
# auth_scopes:
|
||||
# "${VIEWER_TOKEN}": [read]
|
||||
auth_scopes: Optional[Dict[str, List[str]]] = None
|
||||
auth_section = gw_data.get("auth")
|
||||
if isinstance(auth_section, dict) and isinstance(auth_section.get("tokens"), list):
|
||||
auth_scopes = {}
|
||||
for entry in auth_section["tokens"]:
|
||||
if isinstance(entry, dict) and entry.get("token"):
|
||||
auth_scopes[str(entry["token"])] = list(entry.get("scopes", []))
|
||||
elif isinstance(gw_data.get("auth_scopes"), dict):
|
||||
auth_scopes = {
|
||||
str(tok): list(scopes)
|
||||
for tok, scopes in gw_data["auth_scopes"].items()
|
||||
}
|
||||
# Either shape may add ``agents: [...]`` per token (Issue #6253):
|
||||
# auth_scopes:
|
||||
# "${SUPPORT_TOKEN}": {scopes: [read, write], agents: [support]}
|
||||
auth_scopes = _parse_auth_scopes(gw_data)
|
||||
|
||||
raw_trusted = gw_data.get("trusted_proxies", [])
|
||||
trusted_proxies = (
|
||||
|
||||
@@ -3841,12 +3841,19 @@ class ConversationRequestProtocol(Protocol):
|
||||
# import lives in core.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
GatewayRequestKind = Literal["approval", "choice", "input"]
|
||||
GatewayRequestKind = Literal["approval", "choice", "input", "secret"]
|
||||
"""Closed set of server→client interactive request kinds.
|
||||
|
||||
* ``approval`` — allow/deny; the reply ``value`` is ``"allow"`` or ``"deny"``.
|
||||
* ``choice`` — select one of ``options``; the reply ``value`` is the option.
|
||||
* ``input`` — free text; the reply ``value`` is the entered text.
|
||||
* ``secret`` — a credential (password, API key, one-time code) captured
|
||||
mid-turn (Issue #6223). Clients MUST render a masked input. Isolation
|
||||
contract: the reply ``value`` is returned *only* to the caller of
|
||||
:meth:`GatewayRequestChannelProtocol.request`; implementations MUST NOT log,
|
||||
journal, echo, or persist it, and the caller (a tool/callback) MUST NOT
|
||||
return it to the model or append it to history. A tool obtains one via
|
||||
``get_request_channel().request(GatewayServerRequest(kind="secret", ...))``.
|
||||
"""
|
||||
|
||||
|
||||
@@ -3893,11 +3900,13 @@ class GatewayServerReply:
|
||||
Attributes:
|
||||
request_id: The id of the request being answered.
|
||||
value: ``"allow"``/``"deny"`` for approval, the chosen option for
|
||||
choice, or the entered text for input.
|
||||
choice, the entered text for input, or the credential for secret.
|
||||
Excluded from ``repr`` so a logged/traced reply never leaks a
|
||||
secret.
|
||||
"""
|
||||
|
||||
request_id: str
|
||||
value: str
|
||||
value: str = field(repr=False)
|
||||
|
||||
def as_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to a serializable dictionary."""
|
||||
|
||||
@@ -7846,11 +7846,24 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
'task_name', 'task_description', 'task_id',
|
||||
'execute_tool_fn', 'stream_callback', 'emit_events',
|
||||
'console', 'reasoning_steps',
|
||||
'reasoning_effort', 'thinking_budget', # translated below
|
||||
}
|
||||
for k, v in kwargs.items():
|
||||
if k not in _internal and v is not None:
|
||||
params[k] = v
|
||||
|
||||
# ── Reasoning effort ────────────────────────────────────────────
|
||||
# Translate the unified level (Issue #4452) like
|
||||
# _build_completion_params does, but in the Responses API shape
|
||||
# ``reasoning: {"effort": ...}``. A per-call value wins over the
|
||||
# instance setting; an explicit ``reasoning`` object is kept as is.
|
||||
effort = kwargs.get('reasoning_effort', getattr(self, 'reasoning_effort', None))
|
||||
if effort is not None and 'reasoning' not in params:
|
||||
from ..thinking.effort import resolve_reasoning_params
|
||||
native = resolve_reasoning_params(effort, params['model'])
|
||||
if 'reasoning_effort' in native:
|
||||
params['reasoning'] = {'effort': native['reasoning_effort']}
|
||||
|
||||
# Match Chat Completions after merging overrides: reasoning models
|
||||
# must not regain unsupported sampling parameters through either path.
|
||||
from .model_capabilities import is_reasoning_model
|
||||
|
||||
@@ -377,6 +377,15 @@ def process_stream_chunks(chunks):
|
||||
print(f"Error processing chunks: {e}")
|
||||
return None
|
||||
|
||||
def _env_timeout() -> Optional[float]:
|
||||
"""Return the OPENAI_TIMEOUT env var as seconds, or None when unset/invalid."""
|
||||
value = os.environ.get("OPENAI_TIMEOUT")
|
||||
try:
|
||||
return float(value) if value else None
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
class OpenAIClient:
|
||||
"""
|
||||
Unified OpenAI client wrapper for sync/async operations.
|
||||
@@ -411,6 +420,8 @@ class OpenAIClient:
|
||||
except (TypeError, ValueError):
|
||||
max_retries = None
|
||||
self.max_retries = max_retries
|
||||
# Per-request timeout (seconds) from OPENAI_TIMEOUT; unset keeps the SDK default (600s).
|
||||
self.timeout = _env_timeout()
|
||||
|
||||
# For local servers like LM Studio, allow minimal API key
|
||||
if self.base_url and not self.api_key:
|
||||
@@ -479,6 +490,8 @@ class OpenAIClient:
|
||||
client_kwargs = {"api_key": self.api_key, "base_url": self.base_url}
|
||||
if self.max_retries is not None:
|
||||
client_kwargs["max_retries"] = self.max_retries
|
||||
if getattr(self, "timeout", None) is not None:
|
||||
client_kwargs["timeout"] = self.timeout
|
||||
self._sync_client = OpenAI(**client_kwargs)
|
||||
return self._sync_client
|
||||
|
||||
@@ -521,6 +534,8 @@ class OpenAIClient:
|
||||
client_kwargs = {"api_key": self.api_key, "base_url": self.base_url}
|
||||
if self.max_retries is not None:
|
||||
client_kwargs["max_retries"] = self.max_retries
|
||||
if getattr(self, "timeout", None) is not None:
|
||||
client_kwargs["timeout"] = self.timeout
|
||||
# Use the plain default httpx client rather than openai's
|
||||
# internal wrapper: the wrapper's ``__del__`` schedules
|
||||
# ``aclose()`` on whichever loop is running at GC time. When a
|
||||
@@ -938,6 +953,12 @@ class OpenAIClient:
|
||||
if not is_reasoning_model(model):
|
||||
params["temperature"] = temperature
|
||||
|
||||
# Chat Completions takes ``reasoning_effort``; the Responses API takes
|
||||
# the same level as ``reasoning: {"effort": ...}``.
|
||||
reasoning_effort = kwargs.get("reasoning_effort")
|
||||
if reasoning_effort:
|
||||
params["reasoning"] = {"effort": reasoning_effort}
|
||||
|
||||
return params
|
||||
|
||||
@staticmethod
|
||||
@@ -3080,7 +3101,7 @@ def get_openai_client(api_key: Optional[str] = None, base_url: Optional[str] = N
|
||||
normalized_max_retries = int(_env_retries)
|
||||
except (TypeError, ValueError):
|
||||
normalized_max_retries = None
|
||||
current_params = (normalized_api_key, normalized_base_url, normalized_max_retries)
|
||||
current_params = (normalized_api_key, normalized_base_url, normalized_max_retries, _env_timeout())
|
||||
|
||||
# Thread-safe client creation
|
||||
with _global_client_lock:
|
||||
|
||||
@@ -572,6 +572,14 @@ class FileMemory:
|
||||
promoted.append(item.id)
|
||||
|
||||
if promoted:
|
||||
# Promotion writes the same store as add_long_term, so retain the
|
||||
# same highest-importance, newest-first items within its capacity.
|
||||
limit = self.config["long_term_limit"]
|
||||
if len(self._long_term) > limit:
|
||||
self._long_term.sort(
|
||||
key=lambda item: (item.importance, item.created_at), reverse=True
|
||||
)
|
||||
self._long_term = self._long_term[:limit]
|
||||
self._save_long_term()
|
||||
self._log(f"Auto-promoted {len(promoted)} items to long-term memory")
|
||||
|
||||
|
||||
@@ -104,7 +104,12 @@ class PlanStorage:
|
||||
if filename is None:
|
||||
filename = self._get_plan_filename(plan)
|
||||
|
||||
path = os.path.join(self.plans_dir, filename)
|
||||
from ..utils.atomic_io import safe_join
|
||||
# Keep ".md" as the suffix so listing/lookup still find the file.
|
||||
stem, ext = os.path.splitext(filename)
|
||||
if ext.lower() != ".md":
|
||||
stem, ext = filename, ".md"
|
||||
path = safe_join(self.plans_dir, stem, ext)
|
||||
|
||||
# Convert plan to markdown
|
||||
markdown = plan.to_markdown()
|
||||
@@ -215,6 +220,29 @@ class PlanStorage:
|
||||
|
||||
return False
|
||||
|
||||
def _todo_path(self, name: str) -> str:
|
||||
"""Confined path for a todo name, reusing a pre-existing legacy file.
|
||||
|
||||
Legacy files (``<name>.json`` written before Issue #6238) are reused
|
||||
only when ``name`` is a plain filename, so upgrades keep working
|
||||
without ever allowing traversal out of ``todos_dir``.
|
||||
"""
|
||||
from ..utils.atomic_io import safe_join
|
||||
path = safe_join(self.todos_dir, name, ".json")
|
||||
name = str(name)
|
||||
if (
|
||||
not os.path.exists(path)
|
||||
and name not in ("", ".", "..")
|
||||
and "\0" not in name
|
||||
and "%" not in name
|
||||
and os.path.basename(name) == name
|
||||
and not (os.altsep and os.altsep in name)
|
||||
):
|
||||
legacy = os.path.join(self.todos_dir, f"{name}.json")
|
||||
if os.path.isfile(legacy):
|
||||
return legacy
|
||||
return path
|
||||
|
||||
def save_todo(self, todo: TodoList, name: str = "current") -> str:
|
||||
"""
|
||||
Save a todo list to storage.
|
||||
@@ -226,11 +254,10 @@ class PlanStorage:
|
||||
Returns:
|
||||
Path to the saved file
|
||||
"""
|
||||
filename = f"{name}.json"
|
||||
path = os.path.join(self.todos_dir, filename)
|
||||
from ..utils.atomic_io import atomic_write_text
|
||||
path = self._todo_path(name)
|
||||
|
||||
# Atomic write so a crash mid-write never truncates the previous todos.
|
||||
from ..utils.atomic_io import atomic_write_text
|
||||
atomic_write_text(path, todo.to_json())
|
||||
|
||||
logger.debug(f"Saved todo list to {path}")
|
||||
@@ -246,8 +273,7 @@ class PlanStorage:
|
||||
Returns:
|
||||
TodoList instance, or None if not found
|
||||
"""
|
||||
filename = f"{name}.json"
|
||||
path = os.path.join(self.todos_dir, filename)
|
||||
path = self._todo_path(name)
|
||||
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
@@ -265,12 +291,13 @@ class PlanStorage:
|
||||
Returns:
|
||||
List of todo list names
|
||||
"""
|
||||
from urllib.parse import unquote
|
||||
todos = []
|
||||
|
||||
for filename in os.listdir(self.todos_dir):
|
||||
if filename.endswith(".json"):
|
||||
name = filename[:-5] # Remove .json extension
|
||||
todos.append(name)
|
||||
# Decode so listed names round-trip through load/delete_todo.
|
||||
todos.append(unquote(filename[:-5]))
|
||||
|
||||
return todos
|
||||
|
||||
@@ -284,8 +311,7 @@ class PlanStorage:
|
||||
Returns:
|
||||
True if deleted, False if not found
|
||||
"""
|
||||
filename = f"{name}.json"
|
||||
path = os.path.join(self.todos_dir, filename)
|
||||
path = self._todo_path(name)
|
||||
|
||||
if os.path.exists(path):
|
||||
os.remove(path)
|
||||
|
||||
@@ -192,6 +192,12 @@ def __getattr__(name: str):
|
||||
_module_cache[name] = workspace_id
|
||||
return workspace_id
|
||||
|
||||
# Session manifest: plane refs + run-in-flight marker (Issue #6248)
|
||||
if name in ("FileSessionManifest", "ManifestEntry", "SessionManifestProtocol"):
|
||||
from . import manifest as _manifest
|
||||
_module_cache[name] = getattr(_manifest, name)
|
||||
return _module_cache[name]
|
||||
|
||||
# Cross-context continuation prompt assembler (session handoff)
|
||||
if name == "build_handoff_prompt":
|
||||
from .handoff import build_handoff_prompt
|
||||
@@ -249,6 +255,9 @@ __all__ = [
|
||||
"clear_session_context",
|
||||
"workspace_id",
|
||||
"build_handoff_prompt",
|
||||
"FileSessionManifest",
|
||||
"ManifestEntry",
|
||||
"SessionManifestProtocol",
|
||||
"MessageOrigin",
|
||||
"wrap_inter_agent",
|
||||
# Encryption at rest
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
"""
|
||||
Session manifest: one durable record per ``session_id`` holding a write-ahead
|
||||
"run in flight" marker (Issue #6248), so a later launch can tell that a run on
|
||||
this session died before reaching a terminal outcome.
|
||||
|
||||
Usage::
|
||||
|
||||
from praisonaiagents.session.manifest import FileSessionManifest
|
||||
|
||||
manifest = FileSessionManifest()
|
||||
entry = manifest.begin_run("my-session", "run-42", run_input="fix the bug")
|
||||
... # run the agent
|
||||
manifest.end_run("my-session", "completed", run_id="run-42")
|
||||
|
||||
# Next launch:
|
||||
stale = manifest.interrupted("my-session") # entry or None
|
||||
|
||||
The run input itself is never written to disk: only a SHA-256 digest of it
|
||||
is kept (``run_digest``), enough to recognise a retry of the same task without
|
||||
storing prompts (which may contain secrets) in plaintext.
|
||||
|
||||
A run that keeps dying is bounded by ``max_attempts``: once the budget is
|
||||
spent, ``begin_run`` logs the stale run as abandoned and restarts the attempt
|
||||
counter instead of counting it up forever.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterator, Optional, Protocol, runtime_checkable
|
||||
|
||||
from .store import FileLock, _safe_session_filename
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_MAX_ATTEMPTS = 3
|
||||
|
||||
|
||||
@dataclass
|
||||
class ManifestEntry:
|
||||
"""A session's last run state."""
|
||||
|
||||
session_id: str
|
||||
run_id: Optional[str] = None
|
||||
run_status: Optional[str] = None
|
||||
run_in_flight: bool = False
|
||||
run_digest: Optional[str] = None
|
||||
attempt: int = 0
|
||||
updated_at: float = 0.0
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "ManifestEntry":
|
||||
return cls(**{k: data[k] for k in cls.__dataclass_fields__ if k in data})
|
||||
|
||||
|
||||
def _digest(run_input: Optional[str]) -> Optional[str]:
|
||||
if run_input is None:
|
||||
return None
|
||||
return hashlib.sha256(run_input.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SessionManifestProtocol(Protocol):
|
||||
"""Per-session run-in-flight marker."""
|
||||
|
||||
def get(self, session_id: str) -> Optional[ManifestEntry]: ...
|
||||
|
||||
def begin_run(
|
||||
self,
|
||||
session_id: str,
|
||||
run_id: str,
|
||||
run_input: Optional[str] = None,
|
||||
max_attempts: int = DEFAULT_MAX_ATTEMPTS,
|
||||
) -> ManifestEntry: ...
|
||||
|
||||
def end_run(
|
||||
self, session_id: str, status: str, run_id: Optional[str] = None
|
||||
) -> Optional[ManifestEntry]: ...
|
||||
|
||||
def interrupted(self, session_id: str) -> Optional[ManifestEntry]: ...
|
||||
|
||||
|
||||
class FileSessionManifest:
|
||||
"""Default JSON-file manifest: ``<dir>/<session_id>.json``, atomic writes."""
|
||||
|
||||
def __init__(self, directory: Optional[str] = None):
|
||||
"""``directory`` defaults to ``<sessions>/manifests``; callers whose
|
||||
session store is scoped (e.g. per project) should pass the matching
|
||||
directory so equal ids in different scopes never share a manifest."""
|
||||
if directory is None:
|
||||
from ..paths import get_sessions_dir
|
||||
|
||||
directory = str(get_sessions_dir() / "manifests")
|
||||
self._dir = Path(directory)
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def _path(self, session_id: str) -> Path:
|
||||
if not session_id:
|
||||
raise ValueError("session_id is required")
|
||||
# Same injective naming as the conversation store: ids that differ
|
||||
# only in ``:``/``.``/``/`` never share a file, and none escape ``_dir``.
|
||||
return self._dir / f"{_safe_session_filename(session_id)}.json"
|
||||
|
||||
@contextmanager
|
||||
def _locked(self, session_id: str) -> Iterator[None]:
|
||||
"""Serialise read-modify-write across threads *and* processes."""
|
||||
with self._lock, FileLock(str(self._path(session_id))):
|
||||
yield
|
||||
|
||||
def _read(self, session_id: str) -> Optional[ManifestEntry]:
|
||||
path = self._path(session_id)
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as fh:
|
||||
data = json.load(fh)
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
except (OSError, ValueError) as exc:
|
||||
logger.warning("Unreadable session manifest %s: %s", path, exc)
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
data["session_id"] = session_id
|
||||
return ManifestEntry.from_dict(data)
|
||||
|
||||
def _write(self, entry: ManifestEntry) -> ManifestEntry:
|
||||
entry.updated_at = time.time()
|
||||
path = self._path(entry.session_id)
|
||||
self._dir.mkdir(parents=True, exist_ok=True)
|
||||
fd, tmp = tempfile.mkstemp(dir=str(self._dir), suffix=".tmp")
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as fh:
|
||||
json.dump(asdict(entry), fh, default=str)
|
||||
fh.flush()
|
||||
os.fsync(fh.fileno())
|
||||
os.replace(tmp, path)
|
||||
except BaseException:
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
return entry
|
||||
|
||||
def get(self, session_id: str) -> Optional[ManifestEntry]:
|
||||
with self._lock:
|
||||
return self._read(session_id)
|
||||
|
||||
def begin_run(
|
||||
self,
|
||||
session_id: str,
|
||||
run_id: str,
|
||||
run_input: Optional[str] = None,
|
||||
max_attempts: int = DEFAULT_MAX_ATTEMPTS,
|
||||
) -> ManifestEntry:
|
||||
"""Commit the write-ahead in-flight marker for a run on this session.
|
||||
|
||||
Only a digest of ``run_input`` is persisted, never the text itself.
|
||||
If the previous run never released its marker it was interrupted; a new
|
||||
run with the same ``run_input`` counts as another attempt of it. Once
|
||||
``max_attempts`` is exceeded the stale run is abandoned and the counter
|
||||
restarts, so a run that keeps dying is not counted up forever.
|
||||
"""
|
||||
with self._locked(session_id):
|
||||
entry = self._read(session_id) or ManifestEntry(session_id=session_id)
|
||||
digest = _digest(run_input)
|
||||
if entry.run_in_flight and entry.run_digest == digest:
|
||||
attempt = entry.attempt + 1
|
||||
if attempt > max_attempts:
|
||||
logger.warning(
|
||||
"Session %s: run %s abandoned after %d attempts",
|
||||
session_id, entry.run_id, entry.attempt,
|
||||
)
|
||||
attempt = 1
|
||||
else:
|
||||
attempt = 1
|
||||
entry.run_id = run_id
|
||||
entry.run_digest = digest
|
||||
entry.run_status = "running"
|
||||
entry.run_in_flight = True
|
||||
entry.attempt = attempt
|
||||
return self._write(entry)
|
||||
|
||||
def end_run(
|
||||
self, session_id: str, status: str, run_id: Optional[str] = None
|
||||
) -> Optional[ManifestEntry]:
|
||||
"""Release the in-flight marker with a terminal ``status``.
|
||||
|
||||
With ``run_id``, only that run's own marker is released: if another
|
||||
run on the same session has since taken the marker, it is left alone
|
||||
(returns ``None``) so that run's later crash is still reported.
|
||||
"""
|
||||
with self._locked(session_id):
|
||||
entry = self._read(session_id)
|
||||
if entry is None or (run_id is not None and entry.run_id != run_id):
|
||||
return None
|
||||
entry.run_in_flight = False
|
||||
entry.run_status = status
|
||||
return self._write(entry)
|
||||
|
||||
def interrupted(self, session_id: str) -> Optional[ManifestEntry]:
|
||||
"""Return the entry if its last run never reached a terminal outcome."""
|
||||
entry = self.get(session_id)
|
||||
return entry if entry is not None and entry.run_in_flight else None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_MAX_ATTEMPTS",
|
||||
"FileSessionManifest",
|
||||
"ManifestEntry",
|
||||
"SessionManifestProtocol",
|
||||
]
|
||||
@@ -73,28 +73,15 @@ def generate_title(
|
||||
fallback_title = _create_fallback_title(user_msg, max_length)
|
||||
|
||||
try:
|
||||
# Check if we're already in an event loop
|
||||
asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
# No running loop - safe to use asyncio.run
|
||||
try:
|
||||
return asyncio.run(generate_title_async(
|
||||
user_msg, assistant_msg, llm_model, timeout, max_length,
|
||||
primary_model=primary_model,
|
||||
))
|
||||
except Exception:
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.debug("Title generation failed in sync context", exc_info=True)
|
||||
return fallback_title
|
||||
else:
|
||||
# Already in event loop - cannot use asyncio.run
|
||||
# Works with or without a running event loop (Issue #6238).
|
||||
from ..utils.async_bridge import run_coroutine_from_any_context
|
||||
return run_coroutine_from_any_context(generate_title_async(
|
||||
user_msg, assistant_msg, llm_model, timeout, max_length,
|
||||
primary_model=primary_model,
|
||||
), timeout=None)
|
||||
except Exception:
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning(
|
||||
"generate_title() called from running event loop; "
|
||||
"use generate_title_async() instead. Returning fallback title."
|
||||
)
|
||||
logging.getLogger(__name__).debug("Title generation failed", exc_info=True)
|
||||
return fallback_title
|
||||
|
||||
|
||||
|
||||
@@ -645,9 +645,8 @@ def crawl4ai_sync(
|
||||
Returns:
|
||||
Dict with crawl results
|
||||
"""
|
||||
return asyncio.run(
|
||||
crawl4ai(url, css_selector, js_code, wait_for)
|
||||
)
|
||||
from ..utils.async_bridge import run_coroutine_from_any_context
|
||||
return run_coroutine_from_any_context(crawl4ai(url, css_selector, js_code, wait_for), timeout=None)
|
||||
|
||||
def crawl4ai_extract_sync(
|
||||
url: str,
|
||||
@@ -664,9 +663,8 @@ def crawl4ai_extract_sync(
|
||||
Returns:
|
||||
Dict with extracted data
|
||||
"""
|
||||
return asyncio.run(
|
||||
crawl4ai_extract(url, schema, js_code)
|
||||
)
|
||||
from ..utils.async_bridge import run_coroutine_from_any_context
|
||||
return run_coroutine_from_any_context(crawl4ai_extract(url, schema, js_code), timeout=None)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example usage
|
||||
|
||||
@@ -126,11 +126,31 @@ def _matches_glob(fname: str, rel: str, glob: Optional[str]) -> bool:
|
||||
if not glob:
|
||||
return True
|
||||
posix = rel.replace(os.sep, "/")
|
||||
return (
|
||||
if (
|
||||
fnmatch.fnmatch(fname, glob)
|
||||
or fnmatch.fnmatch(rel, glob)
|
||||
or fnmatch.fnmatch(posix, glob)
|
||||
)
|
||||
):
|
||||
return True
|
||||
|
||||
# fnmatch treats "**/" as requiring a slash. Recursive glob segments
|
||||
# also match zero directories, as they do on the ripgrep path.
|
||||
segments = glob.split("/")
|
||||
if "**" not in segments:
|
||||
return False
|
||||
parts = posix.split("/")
|
||||
positions = {0}
|
||||
for segment in segments:
|
||||
if not positions:
|
||||
return False
|
||||
if segment == "**":
|
||||
positions = set(range(min(positions), len(parts) + 1))
|
||||
else:
|
||||
positions = {
|
||||
i + 1 for i in positions
|
||||
if i < len(parts) and fnmatch.fnmatch(parts[i], segment)
|
||||
}
|
||||
return len(parts) in positions
|
||||
|
||||
|
||||
def _iter_files(root: str, glob: Optional[str], ignore_matcher):
|
||||
|
||||
@@ -9,8 +9,10 @@ the previous good file: the old bytes stay intact and ``os.replace`` swaps in th
|
||||
new content atomically on both POSIX and Windows.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
@@ -21,9 +23,52 @@ __all__ = [
|
||||
"atomic_write_json",
|
||||
"atomic_write_text",
|
||||
"quarantine_corrupt_file",
|
||||
"safe_join",
|
||||
"update_json",
|
||||
]
|
||||
|
||||
# Max UTF-8 bytes for the final basename; leaves room under the common
|
||||
# 255-byte limit for ``atomic_write_text``'s ``<basename>.XXXXXXXX.tmp`` file.
|
||||
_MAX_NAME_BYTES = 200
|
||||
_HASHED = re.compile(r"(.*)%~([0-9a-f]{16})", re.DOTALL)
|
||||
_ESCAPE = re.compile(r'[%/\\:*?"<>|\x00-\x1f\x7f]')
|
||||
|
||||
|
||||
def _escape(name: str) -> str:
|
||||
stem = _ESCAPE.sub(lambda m: "%%%02X" % ord(m.group()), name)
|
||||
return "%2E" + stem[1:] if stem.startswith(".") else stem
|
||||
|
||||
|
||||
def safe_join(base: str, name: str, suffix: str = "") -> str:
|
||||
"""Join a caller-supplied ``name`` under ``base`` without escaping it.
|
||||
|
||||
Only ``%``, path separators, control and Windows-reserved characters (and
|
||||
a leading ``.``) are percent-encoded, so ordinary names - including spaces
|
||||
and non-ASCII - keep their pre-#6238 filename, ``..`` and separators can
|
||||
never leave ``base``, and ``urllib.parse.unquote`` recovers the name.
|
||||
Names whose filename would exceed ``_MAX_NAME_BYTES`` are truncated and
|
||||
tagged ``%~<sha256[:16]>``; that listed (unquoted) form maps back to the
|
||||
same file, while no ordinary name can produce a raw ``%~`` (Issue #6238).
|
||||
"""
|
||||
name = str(name)
|
||||
hashed = _HASHED.fullmatch(name)
|
||||
if hashed:
|
||||
stem = _escape(hashed.group(1)) + "%~" + hashed.group(2)
|
||||
else:
|
||||
stem = _escape(name)
|
||||
budget = _MAX_NAME_BYTES - len(suffix.encode("utf-8"))
|
||||
if len(stem.encode("utf-8", "surrogatepass")) > budget:
|
||||
tag = "%~" + hashlib.sha256(name.encode("utf-8", "surrogatepass")).hexdigest()[:16]
|
||||
kept, used = [], len(tag)
|
||||
for tok in re.findall(r"%[0-9A-F]{2}|.", _escape(name), re.DOTALL):
|
||||
used += len(tok.encode("utf-8", "surrogatepass"))
|
||||
if used > budget:
|
||||
break
|
||||
kept.append(tok)
|
||||
stem = "".join(kept) + tag
|
||||
return os.path.join(base, stem + suffix)
|
||||
|
||||
|
||||
# Per-path in-process mutexes so threads in one process serialise before they
|
||||
# even reach the cross-process ``FileLock``. Keyed by absolute path.
|
||||
_path_mutexes: "defaultdict[str, threading.RLock]" = defaultdict(threading.RLock)
|
||||
|
||||
@@ -7,6 +7,7 @@ after it.
|
||||
|
||||
flow = AgentFlow(steps=[...], cache=True) # in-memory, this run
|
||||
flow = AgentFlow(steps=[...], cache=InMemoryStepCache(max_entries=256))
|
||||
flow = AgentFlow(steps=[...], cache=FileStepCache("run.json")) # durable resume
|
||||
|
||||
The key is (step identity, input, previous output, the variables the step can
|
||||
see). Two different steps, or the same step on different input, never collide.
|
||||
@@ -20,11 +21,15 @@ codebase has been full of. Nothing caches unless asked.
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
__all__ = ["StepCacheProtocol", "InMemoryStepCache", "make_step_key", "resolve_step_cache"]
|
||||
__all__ = [
|
||||
"StepCacheProtocol", "InMemoryStepCache", "FileStepCache",
|
||||
"make_step_key", "resolve_step_cache",
|
||||
]
|
||||
|
||||
|
||||
def _snapshot(value: Any) -> Any:
|
||||
@@ -51,6 +56,106 @@ def _identify(step: Any) -> str:
|
||||
return f"{type(step).__name__}:{id(step)}"
|
||||
|
||||
|
||||
_SCALARS = (str, int, float, complex, bool, bytes, type(None), type(Ellipsis))
|
||||
|
||||
|
||||
def _stable(value: Any, seen: set) -> str:
|
||||
# A text form that is the same in every process. repr() is not: a set or
|
||||
# frozenset (``x in {"a", "b"}`` compiles to a frozenset constant) iterates
|
||||
# in an order that depends on the per-process string hash seed.
|
||||
if hasattr(value, "co_code"):
|
||||
return _code_fingerprint(value, seen)
|
||||
if isinstance(value, (set, frozenset)):
|
||||
return "{" + ",".join(sorted(_stable(v, seen) for v in value)) + "}"
|
||||
if isinstance(value, (tuple, list)):
|
||||
return type(value).__name__ + "(" + ",".join(_stable(v, seen) for v in value) + ")"
|
||||
return repr(value)
|
||||
|
||||
|
||||
def _code_fingerprint(code: Any, seen: set) -> str:
|
||||
# Bytecode, names and constants, recursing into nested code objects (whose
|
||||
# repr carries a memory address that would differ in every process).
|
||||
parts = [code.co_code.hex(), repr(code.co_names)]
|
||||
parts.extend(_stable(const, seen) for const in code.co_consts)
|
||||
return hashlib.sha256("|".join(parts).encode("utf-8")).hexdigest()[:16]
|
||||
|
||||
|
||||
def _immutable(value: Any) -> bool:
|
||||
if isinstance(value, _SCALARS):
|
||||
return True
|
||||
if isinstance(value, (tuple, frozenset)):
|
||||
return all(_immutable(v) for v in value)
|
||||
return False
|
||||
|
||||
|
||||
def _function_fingerprint(fn: Any, seen: set) -> str:
|
||||
# The handler's code plus what it was configured with: immutable defaults
|
||||
# and closure values (a factory's ``output="v2"``), and the module-level
|
||||
# Python functions it calls. Mutable captured state (a list it appends to,
|
||||
# a client object) is left out, or the key would change as the run goes.
|
||||
if id(fn) in seen:
|
||||
return "recursive"
|
||||
seen.add(id(fn))
|
||||
parts = [_code_fingerprint(fn.__code__, seen)]
|
||||
captured = list(getattr(fn, "__defaults__", None) or ())
|
||||
captured += sorted((getattr(fn, "__kwdefaults__", None) or {}).items())
|
||||
for cell in getattr(fn, "__closure__", None) or ():
|
||||
try:
|
||||
captured.append(cell.cell_contents)
|
||||
except ValueError: # empty cell
|
||||
continue
|
||||
for value in captured:
|
||||
if hasattr(value, "__code__"):
|
||||
parts.append(_function_fingerprint(value, seen))
|
||||
elif _immutable(value):
|
||||
parts.append(_stable(value, seen))
|
||||
fn_globals = getattr(fn, "__globals__", None) or {}
|
||||
for name in fn.__code__.co_names:
|
||||
helper = fn_globals.get(name)
|
||||
if hasattr(helper, "__code__") and not isinstance(helper, type):
|
||||
parts.append(name + "=" + _function_fingerprint(helper, seen))
|
||||
return hashlib.sha256("|".join(parts).encode("utf-8")).hexdigest()[:16]
|
||||
|
||||
|
||||
_CONFIG_ATTRS = ("role", "goal", "backstory", "instructions", "llm", "action",
|
||||
"description", "expected_output")
|
||||
|
||||
|
||||
def _config(obj: Any) -> Dict[str, Any]:
|
||||
config: Dict[str, Any] = {
|
||||
a: getattr(obj, a) for a in _CONFIG_ATTRS if isinstance(getattr(obj, a, None), str)
|
||||
}
|
||||
tools = getattr(obj, "tools", None)
|
||||
if isinstance(tools, (list, tuple)):
|
||||
config["tools"] = sorted(
|
||||
t if isinstance(t, str) else str(getattr(t, "name", None) or getattr(t, "__name__", type(t).__name__))
|
||||
for t in tools
|
||||
)
|
||||
return config
|
||||
|
||||
|
||||
def _definition(step: Any) -> Optional[str]:
|
||||
# A step edited but kept under the same name must not be served its old
|
||||
# result from a durable cache, so its definition is part of the key: a
|
||||
# handler's code and configuration, or an agent's instructions, model and
|
||||
# tool names.
|
||||
fn = step if hasattr(step, "__code__") else getattr(step, "handler", None)
|
||||
if hasattr(fn, "__code__"):
|
||||
return _function_fingerprint(fn, set())
|
||||
config = _config(step)
|
||||
# AgentFlow wraps an Agent step in a Task(agent=...), so the agent's own
|
||||
# role, instructions and model live one level down.
|
||||
agent = getattr(step, "agent", None)
|
||||
if agent is not None and not isinstance(agent, str):
|
||||
config["agent"] = _config(agent)
|
||||
agent_config = getattr(step, "agent_config", None)
|
||||
if isinstance(agent_config, dict):
|
||||
config["agent_config"] = {k: str(v) for k, v in sorted(agent_config.items())}
|
||||
if not config:
|
||||
return None
|
||||
return hashlib.sha256(json.dumps(config, sort_keys=True).encode("utf-8")).hexdigest()[:16]
|
||||
|
||||
|
||||
def make_step_key(step: Any, previous_output: Any, input_text: str, variables: Dict[str, Any]) -> str:
|
||||
"""A stable key for (this step, these inputs).
|
||||
|
||||
@@ -60,12 +165,13 @@ def make_step_key(step: Any, previous_output: Any, input_text: str, variables: D
|
||||
"""
|
||||
payload = {
|
||||
"step": _identify(step),
|
||||
"definition": _definition(step),
|
||||
"previous": previous_output if isinstance(previous_output, str) else str(previous_output),
|
||||
"input": input_text,
|
||||
# Tag each variable with its type so True and "True" (or 1 and "1") do
|
||||
# not collapse to the same key and serve one call's answer to another.
|
||||
"variables": {
|
||||
k: [type(v).__name__, str(v)]
|
||||
k: [type(v).__name__, _stable(v, set()) if isinstance(v, (set, frozenset)) else str(v)]
|
||||
for k, v in sorted((variables or {}).items())
|
||||
},
|
||||
}
|
||||
@@ -163,6 +269,78 @@ class InMemoryStepCache:
|
||||
return len(self._entries)
|
||||
|
||||
|
||||
class FileStepCache:
|
||||
"""Durable step cache: resume an interrupted flow in a later process.
|
||||
|
||||
flow = AgentFlow(steps=[...], cache=FileStepCache("run.json"))
|
||||
flow.run("input") # crashes at step 4
|
||||
flow.run("input") # steps 1-3 are served from run.json; step 4 runs
|
||||
|
||||
Resume is content-addressed, not index-based: each key covers the step's
|
||||
name and definition (a handler's code, its immutable defaults and closure
|
||||
values and the module-level functions it calls; an agent's role, goal,
|
||||
instructions, model and tool names), the input, the previous step's output
|
||||
and the variables, so renaming or editing a step or changing the input
|
||||
misses the cache from that step on and it runs again. Not covered: a
|
||||
change in mutable captured state, a method or imported library a handler
|
||||
calls, or a tool's implementation. After such an edit, call ``clear()``
|
||||
(or delete the file) to re-run. Steps without a stable name never hit
|
||||
across processes; they re-run.
|
||||
|
||||
Every ``set`` is a locked read-modify-write of the file, so several caches
|
||||
(or processes) on one path add to it rather than erase each other's
|
||||
entries, and a crash mid-write keeps the last good file. A value that JSON
|
||||
cannot store unchanged (an object, int dict keys, a tuple) is rejected (the
|
||||
workflow logs a warning and carries on) rather than stored lossily.
|
||||
"""
|
||||
|
||||
def __init__(self, path: str):
|
||||
self.path = os.path.abspath(os.path.expanduser(str(path)))
|
||||
self._lock = threading.Lock()
|
||||
self._entries: Dict[str, Dict[str, Any]] = {}
|
||||
if os.path.exists(self.path):
|
||||
with open(self.path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError(f"{self.path} is not a step cache file.")
|
||||
self._entries = data
|
||||
|
||||
def get(self, key: str) -> Optional[Dict[str, Any]]:
|
||||
with self._lock:
|
||||
entry = self._entries.get(key)
|
||||
return copy.deepcopy(entry) if entry is not None else None
|
||||
|
||||
def set(self, key: str, value: Dict[str, Any]) -> None:
|
||||
# Round-trip through JSON first so an unserialisable value raises here
|
||||
# and never reaches memory or disk in a form a later run cannot read.
|
||||
snapshot = json.loads(json.dumps(value))
|
||||
if snapshot != value:
|
||||
# e.g. {1: "x"} -> {"1": "x"} or (a, b) -> [a, b]: a resumed run
|
||||
# would see different data than the live run did.
|
||||
raise TypeError("Step cache value does not survive a JSON round trip unchanged.")
|
||||
from ..utils.atomic_io import update_json
|
||||
|
||||
def _add(entries: Any) -> Dict[str, Any]:
|
||||
if not isinstance(entries, dict):
|
||||
raise ValueError(f"{self.path} is not a step cache file.")
|
||||
entries[key] = snapshot
|
||||
return dict(entries)
|
||||
|
||||
with self._lock:
|
||||
os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True)
|
||||
self._entries = update_json(self.path, _add)
|
||||
|
||||
def clear(self) -> None:
|
||||
with self._lock:
|
||||
self._entries = {}
|
||||
if os.path.exists(self.path):
|
||||
os.remove(self.path)
|
||||
|
||||
def __len__(self) -> int:
|
||||
with self._lock:
|
||||
return len(self._entries)
|
||||
|
||||
|
||||
def resolve_step_cache(cache: Any) -> Optional[Any]:
|
||||
"""Turn a ``cache=`` value into a cache object, or None for 'do not cache'."""
|
||||
if cache is None or cache is False:
|
||||
|
||||
@@ -3005,17 +3005,16 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
|
||||
result = agent.generate(prompt)
|
||||
# Extract URL from result and store in variables for next step
|
||||
if hasattr(result, 'data') and result.data:
|
||||
image_url = result.data[0].url if hasattr(result.data[0], 'url') else None
|
||||
image_url = getattr(result.data[0], 'url', None)
|
||||
if image_url:
|
||||
variables['_last_image_url'] = image_url
|
||||
# Save to file if output path specified
|
||||
if output_path:
|
||||
try:
|
||||
import urllib.request
|
||||
urllib.request.urlretrieve(image_url, output_path)
|
||||
variables['_last_image_file'] = output_path
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save image to {output_path}: {e}")
|
||||
# Save to file if output path specified (handles b64 and URL responses)
|
||||
if output_path:
|
||||
try:
|
||||
result.save(output_path)
|
||||
variables['_last_image_file'] = output_path
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save image to {output_path}: {e}")
|
||||
return result
|
||||
|
||||
# OCRAgent - extract with retry logic for external API failures
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.21"
|
||||
version = "1.7.22"
|
||||
description = "Open-source Python AI agent framework that runs on your own computers, even fully offline."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
"""Live check that async chat with a tool returns a model-written final answer.
|
||||
|
||||
Exercises the post-tool follow-up request in ``ChatMixin._achat_completion``
|
||||
by turning off the default unified dispatch, so ``_achat_impl`` takes the
|
||||
legacy OpenAI-client branch. Skipped unless ``OPENAI_API_KEY`` is set.
|
||||
"""
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
|
||||
|
||||
def get_secret_number() -> str:
|
||||
"""Return the secret number."""
|
||||
return "4217"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No API key")
|
||||
async def test_achat_with_tool_returns_final_answer_live():
|
||||
agent = Agent(
|
||||
name="test",
|
||||
instructions="You are a helpful assistant",
|
||||
llm="gpt-4o-mini",
|
||||
tools=[get_secret_number],
|
||||
)
|
||||
# The default unified path never calls _achat_completion.
|
||||
agent._use_unified_llm_dispatch = False
|
||||
result = await agent.astart(
|
||||
"Use the get_secret_number tool, then reply with one sentence that "
|
||||
"contains the secret number."
|
||||
)
|
||||
print(f"Agent output: {result}")
|
||||
assert result
|
||||
assert "4217" in str(result)
|
||||
# Without the fix the raw tool output ("4217") is returned unchanged.
|
||||
assert str(result).strip() != "4217"
|
||||
@@ -169,3 +169,264 @@ async def test_execute_tool_async_resolves_aliased_function_name():
|
||||
)
|
||||
|
||||
assert result == "aliased-ran"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_achat_completion_streams_final_response_after_tool_calls():
|
||||
"""After running tools, the legacy async path must ask the model for a
|
||||
final answer instead of failing and returning the raw tool output.
|
||||
|
||||
``_achat_completion`` used to reference an undefined ``stream`` name and
|
||||
``_get_llm_functions`` (only defined in ``agent.py``), so the follow-up
|
||||
completion always raised ``NameError``, which was swallowed and the raw
|
||||
tool results were returned.
|
||||
"""
|
||||
agent = Agent(name="test", instructions="You are helpful", llm="gpt-4o-mini")
|
||||
|
||||
def lookup(query: str) -> str:
|
||||
return "tool output"
|
||||
|
||||
tool_call = SimpleNamespace(
|
||||
id="call-1",
|
||||
function=SimpleNamespace(name="lookup", arguments='{"query": "x"}'),
|
||||
)
|
||||
response = SimpleNamespace(
|
||||
choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=[tool_call]))]
|
||||
)
|
||||
|
||||
def _chunk(content):
|
||||
return SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(content=content))]
|
||||
)
|
||||
|
||||
async def _stream():
|
||||
for part in ("final ", "answer"):
|
||||
yield _chunk(part)
|
||||
|
||||
create = AsyncMock(return_value=_stream())
|
||||
fake_client = SimpleNamespace(
|
||||
async_client=SimpleNamespace(
|
||||
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
|
||||
)
|
||||
)
|
||||
final = SimpleNamespace(
|
||||
choices=[SimpleNamespace(message=SimpleNamespace(content="final answer"))]
|
||||
)
|
||||
process_stream_chunks = MagicMock(return_value=final)
|
||||
|
||||
with patch.object(
|
||||
type(agent), "_openai_client", new=fake_client
|
||||
), patch.object(
|
||||
agent, "execute_tool_async", new_callable=AsyncMock, return_value="tool output"
|
||||
), patch(
|
||||
"praisonaiagents.agent.agent._get_llm_functions",
|
||||
return_value={"process_stream_chunks": process_stream_chunks},
|
||||
), patch(
|
||||
"praisonaiagents.agent.chat_mixin._get_display_functions"
|
||||
) as display:
|
||||
result = await agent._achat_completion(response, tools=[lookup])
|
||||
|
||||
display.return_value.__getitem__.return_value.assert_not_called()
|
||||
create.assert_awaited_once()
|
||||
assert create.await_args.kwargs["stream"] is True
|
||||
process_stream_chunks.assert_called_once()
|
||||
assert result == "final answer"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_achat_completion_follow_up_keeps_conversation_and_skips_empty_chunks():
|
||||
"""The post-tool follow-up request must keep the original conversation (so
|
||||
the final answer still sees the user's question), and stream chunks with no
|
||||
``choices`` (e.g. Azure content-filter annotations) must be skipped rather
|
||||
than raising ``IndexError`` and falling back to the raw tool output.
|
||||
"""
|
||||
agent = Agent(name="test", instructions="You are helpful", llm="gpt-4o-mini")
|
||||
|
||||
def lookup(query: str) -> str:
|
||||
return "tool output"
|
||||
|
||||
tool_call = SimpleNamespace(
|
||||
id="call-1",
|
||||
function=SimpleNamespace(name="lookup", arguments='{"query": "x"}'),
|
||||
)
|
||||
response = SimpleNamespace(
|
||||
choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=[tool_call]))]
|
||||
)
|
||||
conversation = [
|
||||
{"role": "system", "content": "You are helpful"},
|
||||
{"role": "user", "content": "Look up x and answer in French"},
|
||||
]
|
||||
|
||||
async def _stream():
|
||||
yield SimpleNamespace(choices=[])
|
||||
yield SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(content="réponse"))]
|
||||
)
|
||||
|
||||
create = AsyncMock(return_value=_stream())
|
||||
fake_client = SimpleNamespace(
|
||||
async_client=SimpleNamespace(
|
||||
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
|
||||
)
|
||||
)
|
||||
final = SimpleNamespace(
|
||||
choices=[SimpleNamespace(message=SimpleNamespace(content="réponse"))]
|
||||
)
|
||||
process_stream_chunks = MagicMock(return_value=final)
|
||||
|
||||
with patch.object(
|
||||
type(agent), "_openai_client", new=fake_client
|
||||
), patch.object(
|
||||
agent, "execute_tool_async", new_callable=AsyncMock, return_value="tool output"
|
||||
), patch(
|
||||
"praisonaiagents.agent.agent._get_llm_functions",
|
||||
return_value={"process_stream_chunks": process_stream_chunks},
|
||||
), patch(
|
||||
"praisonaiagents.agent.chat_mixin._get_display_functions"
|
||||
) as display:
|
||||
result = await agent._achat_completion(
|
||||
response, tools=[lookup], messages=conversation
|
||||
)
|
||||
|
||||
display.return_value.__getitem__.return_value.assert_not_called()
|
||||
sent = create.await_args.kwargs["messages"]
|
||||
assert sent[: len(conversation)] == conversation
|
||||
assert "tool output" in sent[-1]["content"]
|
||||
# The caller's list is not mutated.
|
||||
assert len(conversation) == 2
|
||||
assert result == "réponse"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_achat_impl_forwards_conversation_to_tool_follow_up():
|
||||
"""The legacy ``_achat_impl`` tool branch must pass the conversation to
|
||||
``_achat_completion`` so the follow-up request still has the user's prompt.
|
||||
"""
|
||||
agent = Agent(name="test", instructions="You are helpful", llm="gpt-4o-mini")
|
||||
agent._using_custom_llm = False
|
||||
agent._use_unified_llm_dispatch = False
|
||||
|
||||
def lookup(query: str) -> str:
|
||||
return "tool output"
|
||||
|
||||
tool_call = SimpleNamespace(
|
||||
id="call-1",
|
||||
function=SimpleNamespace(name="lookup", arguments='{"query": "x"}'),
|
||||
)
|
||||
tool_response = SimpleNamespace(
|
||||
choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=[tool_call]))]
|
||||
)
|
||||
|
||||
async def _stream():
|
||||
yield SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(content="final answer"))]
|
||||
)
|
||||
|
||||
create = AsyncMock(side_effect=[tool_response, _stream()])
|
||||
fake_client = SimpleNamespace(
|
||||
async_client=SimpleNamespace(
|
||||
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
|
||||
)
|
||||
)
|
||||
final = SimpleNamespace(
|
||||
choices=[SimpleNamespace(message=SimpleNamespace(content="final answer"))]
|
||||
)
|
||||
prompt_messages = [{"role": "user", "content": "Look up x for me"}]
|
||||
|
||||
with patch.object(type(agent), "_openai_client", new=fake_client), patch.object(
|
||||
agent, "execute_tool_async", new_callable=AsyncMock, return_value="tool output"
|
||||
), patch(
|
||||
"praisonaiagents.agent.agent._get_llm_functions",
|
||||
return_value={"process_stream_chunks": MagicMock(return_value=final)},
|
||||
), patch.object(
|
||||
agent, "_atrigger_after_agent_hook", new_callable=AsyncMock, side_effect=lambda p, r, t: r
|
||||
), patch.object(agent, "_execute_callback_and_display"), patch.object(
|
||||
agent, "_build_messages", return_value=(list(prompt_messages), "Look up x for me")
|
||||
), patch.object(agent, "_persist_message"), patch.object(
|
||||
agent, "_format_tools_for_completion", return_value=[{"type": "function"}]
|
||||
):
|
||||
result = await agent._achat_impl(
|
||||
prompt="Look up x for me",
|
||||
temperature=1.0,
|
||||
tools=[lookup],
|
||||
output_json=None,
|
||||
output_pydantic=None,
|
||||
reasoning_steps=False,
|
||||
stream=False,
|
||||
task_name=None,
|
||||
task_description=None,
|
||||
task_id=None,
|
||||
config=None,
|
||||
force_retrieval=False,
|
||||
skip_retrieval=True,
|
||||
attachments=None,
|
||||
_trace_emitter=MagicMock(),
|
||||
)
|
||||
|
||||
assert create.await_count == 2
|
||||
follow_up = create.await_args_list[1].kwargs["messages"]
|
||||
assert {"role": "user", "content": "Look up x for me"} in follow_up
|
||||
assert result == "final answer"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"follow_up_chunk",
|
||||
[
|
||||
pytest.param(SimpleNamespace(choices=[]), id="empty-choices"),
|
||||
pytest.param(
|
||||
SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(content=None))]
|
||||
),
|
||||
id="choice-without-content",
|
||||
),
|
||||
pytest.param(
|
||||
SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(content="\n "))]
|
||||
),
|
||||
id="whitespace-only-content",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_achat_completion_falls_back_when_follow_up_has_no_answer_chunks(follow_up_chunk):
|
||||
"""A follow-up stream with no answer text (only empty-``choices`` chunks,
|
||||
choices without ``delta.content`` or whitespace-only content) must log an error and return the tool
|
||||
results, not an empty string.
|
||||
"""
|
||||
agent = Agent(name="test", instructions="You are helpful", llm="gpt-4o-mini")
|
||||
|
||||
def lookup(query: str) -> str:
|
||||
return "tool output"
|
||||
|
||||
tool_call = SimpleNamespace(
|
||||
id="call-1",
|
||||
function=SimpleNamespace(name="lookup", arguments='{"query": "x"}'),
|
||||
)
|
||||
response = SimpleNamespace(
|
||||
choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=[tool_call]))]
|
||||
)
|
||||
|
||||
async def _stream():
|
||||
yield follow_up_chunk
|
||||
|
||||
create = AsyncMock(return_value=_stream())
|
||||
fake_client = SimpleNamespace(
|
||||
async_client=SimpleNamespace(
|
||||
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
|
||||
)
|
||||
)
|
||||
process_stream_chunks = MagicMock()
|
||||
|
||||
with patch.object(type(agent), "_openai_client", new=fake_client), patch.object(
|
||||
agent, "execute_tool_async", new_callable=AsyncMock, return_value="tool output"
|
||||
), patch(
|
||||
"praisonaiagents.agent.agent._get_llm_functions",
|
||||
return_value={"process_stream_chunks": process_stream_chunks},
|
||||
), patch(
|
||||
"praisonaiagents.agent.chat_mixin._get_display_functions"
|
||||
) as display:
|
||||
result = await agent._achat_completion(response, tools=[lookup])
|
||||
|
||||
display.return_value.__getitem__.return_value.assert_called()
|
||||
process_stream_chunks.assert_not_called()
|
||||
assert result == "tool output"
|
||||
@@ -0,0 +1,113 @@
|
||||
"""The custom-LLM (``provider/model``) path sends each user turn once.
|
||||
|
||||
``chat()`` stores the user turn in ``chat_history`` before the model call, and
|
||||
``LLM._build_messages`` appends ``prompt`` after ``chat_history`` itself, so
|
||||
the stored copy must not be forwarded as history too. It was, on every
|
||||
litellm route (``ollama/``, ``anthropic/``, ``openai/``, ``base_url=``), which
|
||||
billed the prompt twice. ``litellm.completion`` is stubbed; nothing leaves the
|
||||
machine.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("litellm")
|
||||
|
||||
from praisonaiagents import Agent # noqa: E402
|
||||
|
||||
MODEL = "anthropic/claude-3-5-sonnet-20240620"
|
||||
|
||||
|
||||
def _litellm():
|
||||
import litellm
|
||||
return litellm
|
||||
|
||||
|
||||
def _reply(model, text):
|
||||
return _litellm().ModelResponse(model=model, usage={
|
||||
"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
|
||||
choices=[{"index": 0, "finish_reason": "stop",
|
||||
"message": {"role": "assistant", "content": text}}])
|
||||
|
||||
|
||||
def _chunks(model, text):
|
||||
litellm = _litellm()
|
||||
yield litellm.ModelResponseStream(model=model, choices=[{
|
||||
"index": 0, "delta": {"role": "assistant", "content": text}}])
|
||||
yield litellm.ModelResponseStream(model=model, choices=[{
|
||||
"index": 0, "delta": {}, "finish_reason": "stop"}])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sent(monkeypatch):
|
||||
"""Record the ``messages`` of every model call; answer "A1", "A2", ..."""
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
monkeypatch.setattr(Agent, "_default_model", None)
|
||||
monkeypatch.setattr(Agent, "_default_model_checked", False)
|
||||
requests = []
|
||||
|
||||
def fake(*a, **k):
|
||||
requests.append(k.get("messages"))
|
||||
text = f"A{len(requests)}"
|
||||
if k.get("stream"):
|
||||
return _chunks(k.get("model"), text)
|
||||
return _reply(k.get("model"), text)
|
||||
|
||||
async def afake(*a, **k):
|
||||
return fake(*a, **k)
|
||||
|
||||
litellm = _litellm()
|
||||
monkeypatch.setattr(litellm, "completion", fake)
|
||||
monkeypatch.setattr(litellm, "acompletion", afake)
|
||||
return requests
|
||||
|
||||
|
||||
def _turns(messages):
|
||||
return [(m["role"], m["content"]) for m in messages if m["role"] != "system"]
|
||||
|
||||
|
||||
def _agent():
|
||||
return Agent(name="p", instructions="Be brief.", llm=MODEL, output="silent")
|
||||
|
||||
|
||||
def test_chat_sends_the_prompt_once(sent):
|
||||
agent = _agent()
|
||||
agent.chat("Q1")
|
||||
agent.chat("Q2")
|
||||
assert _turns(sent[0]) == [("user", "Q1")]
|
||||
assert _turns(sent[1]) == [("user", "Q1"), ("assistant", "A1"), ("user", "Q2")]
|
||||
|
||||
|
||||
def test_achat_sends_the_prompt_once(sent):
|
||||
agent = _agent()
|
||||
|
||||
async def run():
|
||||
await agent.achat("Q1")
|
||||
await agent.achat("Q2")
|
||||
|
||||
asyncio.run(run())
|
||||
assert _turns(sent[0]) == [("user", "Q1")]
|
||||
assert _turns(sent[1]) == [("user", "Q1"), ("assistant", "A1"), ("user", "Q2")]
|
||||
|
||||
|
||||
def test_streamed_turn_sends_the_prompt_once(sent):
|
||||
agent = _agent()
|
||||
assert "".join(agent.start("Q1", stream=True)) == "A1"
|
||||
assert _turns(sent[0]) == [("user", "Q1")]
|
||||
|
||||
|
||||
def test_history_still_records_each_turn_once(sent):
|
||||
agent = _agent()
|
||||
agent.chat("Q1")
|
||||
agent.chat("Q2")
|
||||
assert [(m["role"], m["content"]) for m in agent.chat_history] == [
|
||||
("user", "Q1"), ("assistant", "A1"), ("user", "Q2"), ("assistant", "A2")]
|
||||
|
||||
|
||||
def test_repeating_the_last_question_is_still_one_turn(sent):
|
||||
# A user who asks the same thing twice in a row sends it once per request,
|
||||
# with the earlier exchange kept as history.
|
||||
agent = _agent()
|
||||
agent.chat("Q1")
|
||||
agent.chat("Q1")
|
||||
assert _turns(sent[1]) == [("user", "Q1"), ("assistant", "A1"), ("user", "Q1")]
|
||||
@@ -38,7 +38,7 @@ def test_cache_hit_refreshes_current_default_tools(agent, monkeypatch):
|
||||
builds = []
|
||||
monkeypatch.setattr(
|
||||
agent, '_resolve_harness_base_prompt',
|
||||
lambda: builds.append('build') or 'stable prefix',
|
||||
lambda tools=None: builds.append('build') or 'stable prefix',
|
||||
)
|
||||
first = agent._build_system_prompt(tools=None)
|
||||
assert 'sample_tool' in first
|
||||
@@ -73,7 +73,7 @@ def test_cache_hit_refreshes_session_context(agent, monkeypatch):
|
||||
def test_cache_hit_retains_prefix_reuse_and_appends_suffix_once(agent, monkeypatch):
|
||||
"""Reuse the stable prefix while adding the suffix exactly once."""
|
||||
calls = []
|
||||
agent._resolve_harness_base_prompt = lambda: calls.append('build') or 'stable prefix'
|
||||
agent._resolve_harness_base_prompt = lambda tools=None: calls.append('build') or 'stable prefix'
|
||||
monkeypatch.setenv('PRAISONAI_APPEND_SYSTEM_PROMPT', 'per-turn suffix')
|
||||
first = agent._build_system_prompt()
|
||||
second = agent._build_system_prompt()
|
||||
|
||||
@@ -0,0 +1,230 @@
|
||||
"""The system prompt carries only guidance that applies to the agent.
|
||||
|
||||
* The model-family file-editing hint ("use edit_file ..." / "use
|
||||
apply_patch ...") is for agents that can edit files; every GPT or Claude
|
||||
agent got it, tools or not.
|
||||
* ``Agent(instructions=...)`` sets both backstory and goal to the
|
||||
instructions, so the prompt repeated them under "Your Goal:".
|
||||
"""
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.tools.edit_tools import apply_patch, edit_file
|
||||
|
||||
OPENAI_HINT = "use edit_file to modify existing files precisely"
|
||||
ANTHROPIC_HINT = "use apply_patch"
|
||||
|
||||
|
||||
def lookup(query: str) -> str:
|
||||
"""Look something up."""
|
||||
return query
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def make(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test-not-used")
|
||||
made = []
|
||||
|
||||
def _make(**kw):
|
||||
kw.setdefault("instructions", "Be brief.")
|
||||
agent = Agent(context=False, **kw)
|
||||
made.append(agent)
|
||||
return agent
|
||||
|
||||
yield _make
|
||||
for agent in made:
|
||||
agent.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tools", [None, [lookup]])
|
||||
def test_agent_without_edit_tools_gets_no_edit_hint(make, tools):
|
||||
agent = make(llm="gpt-4o-mini", tools=tools)
|
||||
assert OPENAI_HINT not in agent._build_system_prompt()
|
||||
|
||||
|
||||
def test_agent_with_edit_file_gets_the_openai_hint(make):
|
||||
agent = make(llm="gpt-4o-mini", tools=[edit_file, lookup])
|
||||
assert OPENAI_HINT in agent._build_system_prompt()
|
||||
|
||||
|
||||
def test_agent_with_apply_patch_gets_the_anthropic_hint(make):
|
||||
agent = make(llm="anthropic/claude-sonnet-4-5", tools=[apply_patch])
|
||||
assert ANTHROPIC_HINT in agent._build_system_prompt()
|
||||
|
||||
|
||||
def test_edit_tool_named_by_string_counts(make):
|
||||
agent = make(llm="gpt-4o-mini")
|
||||
assert OPENAI_HINT in agent._build_system_prompt(tools=["edit_file"])
|
||||
|
||||
|
||||
def test_per_call_tools_decide_the_hint(make):
|
||||
agent = make(llm="gpt-4o-mini", tools=[edit_file])
|
||||
assert OPENAI_HINT not in agent._build_system_prompt(tools=[lookup])
|
||||
assert OPENAI_HINT in agent._build_system_prompt(tools=[edit_file])
|
||||
|
||||
|
||||
def test_cached_prompt_follows_a_change_to_the_agents_tools(make):
|
||||
agent = make(llm="gpt-4o-mini", tools=[lookup])
|
||||
assert OPENAI_HINT not in agent._build_system_prompt()
|
||||
agent.tools = [edit_file]
|
||||
assert OPENAI_HINT in agent._build_system_prompt()
|
||||
agent.tools = [lookup]
|
||||
assert OPENAI_HINT not in agent._build_system_prompt()
|
||||
|
||||
|
||||
def test_explicit_harness_override_is_always_used(make):
|
||||
agent = make(llm="gpt-4o-mini")
|
||||
agent.harness_base_prompt = "House rules."
|
||||
assert agent._build_system_prompt().startswith("House rules.")
|
||||
|
||||
|
||||
def test_instructions_appear_once(make):
|
||||
agent = make(llm="gpt-4o-mini", instructions="Answer in French.")
|
||||
prompt = agent._build_system_prompt()
|
||||
assert prompt.count("Answer in French.") == 1
|
||||
assert "Your Goal:" not in prompt
|
||||
|
||||
|
||||
def test_a_distinct_goal_is_still_stated(make):
|
||||
agent = make(llm="gpt-4o-mini", instructions=None, role="Editor",
|
||||
goal="Tighten the copy", backstory="A veteran sub-editor")
|
||||
prompt = agent._build_system_prompt()
|
||||
assert "Your Goal: Tighten the copy" in prompt
|
||||
assert "A veteran sub-editor" in prompt
|
||||
assert "Your Role: Editor" in prompt
|
||||
|
||||
|
||||
def test_a_tool_renamed_to_edit_file_counts(make):
|
||||
from praisonaiagents import tool
|
||||
|
||||
@tool(name="edit_file")
|
||||
def my_editor(filepath: str, old_string: str, new_string: str) -> str:
|
||||
"""Edit a file."""
|
||||
return "ok"
|
||||
|
||||
agent = make(llm="gpt-4o-mini", tools=[my_editor])
|
||||
assert OPENAI_HINT in agent._build_system_prompt()
|
||||
|
||||
|
||||
def test_nested_tool_list_offers_the_edit_hint(make):
|
||||
# _format_tools_for_completion flattens nested lists of tool dicts.
|
||||
agent = make(llm="gpt-4o-mini", tools=[[EDIT_SCHEMA]])
|
||||
assert OPENAI_HINT in agent._build_system_prompt()
|
||||
|
||||
|
||||
def test_nested_names_the_formatter_drops_get_no_hint(make):
|
||||
# A nested string is not sent to the model, so neither is the hint.
|
||||
agent = make(llm="gpt-4o-mini", tools=[["edit_file"]])
|
||||
assert OPENAI_HINT not in agent._build_system_prompt()
|
||||
|
||||
|
||||
def test_nested_bare_name_dict_gets_no_hint(make):
|
||||
# A nested dict without a function schema is dropped by the formatter.
|
||||
agent = make(llm="gpt-4o-mini", tools=[[{"name": "edit_file"}]])
|
||||
assert agent._format_tools_for_completion() == []
|
||||
assert OPENAI_HINT not in agent._build_system_prompt()
|
||||
|
||||
|
||||
EDIT_SCHEMA = {"type": "function", "function": {
|
||||
"name": "edit_file", "description": "Edit a file.",
|
||||
"parameters": {"type": "object", "properties": {}}}}
|
||||
|
||||
|
||||
def _system_sent(agent, monkeypatch, run):
|
||||
"""Run ``run(agent)`` on the native path; return the system prompt and tool names sent."""
|
||||
from types import SimpleNamespace
|
||||
sent = []
|
||||
sent_tools = []
|
||||
|
||||
def fake(messages=None, temperature=None, tools=None, *a, **k):
|
||||
sent.append(messages if messages is not None else k.get("messages"))
|
||||
# Record the payload the real request builds (inherit, flatten, prune).
|
||||
formatted = agent._format_tools_for_completion(tools)
|
||||
sent_tools.append([t.get("function", {}).get("name") for t in formatted])
|
||||
message = SimpleNamespace(content="ok", tool_calls=None, role="assistant", reasoning_content=None)
|
||||
return SimpleNamespace(choices=[SimpleNamespace(message=message, finish_reason="stop", index=0)],
|
||||
usage=SimpleNamespace(prompt_tokens=1, completion_tokens=1, total_tokens=2),
|
||||
model="gpt-4o-mini")
|
||||
|
||||
async def afake(*a, **k):
|
||||
return fake(*a, **k)
|
||||
|
||||
monkeypatch.setattr(agent, "_chat_completion", fake)
|
||||
monkeypatch.setattr(agent, "_achat_completion_with_retry", afake)
|
||||
run(agent)
|
||||
return next(m["content"] for m in sent[-1] if m["role"] == "system"), sent_tools[-1]
|
||||
|
||||
|
||||
def test_native_chat_follows_per_call_edit_tools(make, monkeypatch):
|
||||
agent = make(llm="gpt-4o-mini")
|
||||
system, tool_names = _system_sent(agent, monkeypatch, lambda a: a.chat("hi", tools=[EDIT_SCHEMA]))
|
||||
assert OPENAI_HINT in system
|
||||
assert "edit_file" in tool_names
|
||||
|
||||
|
||||
def test_nested_edit_schema_is_sent_with_the_hint(make, monkeypatch):
|
||||
agent = make(llm="gpt-4o-mini", tools=[[EDIT_SCHEMA]])
|
||||
system, tool_names = _system_sent(agent, monkeypatch, lambda a: a.chat("hi"))
|
||||
assert OPENAI_HINT in system
|
||||
assert "edit_file" in tool_names
|
||||
|
||||
|
||||
def test_permission_denied_edit_tool_gets_no_hint(make, monkeypatch):
|
||||
agent = make(llm="gpt-4o-mini", tools=[edit_file, lookup])
|
||||
agent._perm_deny = frozenset({"edit_file"})
|
||||
system, tool_names = _system_sent(agent, monkeypatch, lambda a: a.chat("hi"))
|
||||
assert "edit_file" not in tool_names
|
||||
assert OPENAI_HINT not in system
|
||||
|
||||
|
||||
def test_deny_added_after_caching_prunes_cached_tools(make):
|
||||
agent = make(llm="gpt-4o-mini", tools=[edit_file, lookup])
|
||||
assert "edit_file" in [t["function"]["name"] for t in agent._format_tools_for_completion()]
|
||||
agent._perm_deny = frozenset({"edit_file"})
|
||||
assert "edit_file" not in [t["function"]["name"] for t in agent._format_tools_for_completion()]
|
||||
assert OPENAI_HINT not in agent._build_system_prompt()
|
||||
# Lifting the denial restores the tool from the same cache entry.
|
||||
agent._perm_deny = frozenset()
|
||||
assert "edit_file" in [t["function"]["name"] for t in agent._format_tools_for_completion()]
|
||||
|
||||
|
||||
def test_deny_before_first_format_restores_after_lift(make):
|
||||
agent = make(llm="gpt-4o-mini", tools=[edit_file, lookup])
|
||||
agent._perm_deny = frozenset({"edit_file"})
|
||||
assert "edit_file" not in [t["function"]["name"] for t in agent._format_tools_for_completion()]
|
||||
agent._perm_deny = frozenset()
|
||||
assert "edit_file" in [t["function"]["name"] for t in agent._format_tools_for_completion()]
|
||||
|
||||
|
||||
def test_bridge_tool_search_hides_denied_tools(make):
|
||||
agent = make(llm="gpt-4o-mini")
|
||||
agent._tool_search_config = None
|
||||
agent._tool_search_metadata = {"bridge_mode": True, "deferrable_tools": [EDIT_SCHEMA]}
|
||||
agent._perm_deny = frozenset({"edit_file"})
|
||||
out = agent._handle_bridge_tool_call("tool_describe", {"tool_name": "edit_file"})
|
||||
assert "not found" in out
|
||||
agent._perm_deny = frozenset()
|
||||
assert "not found" not in agent._handle_bridge_tool_call("tool_describe", {"tool_name": "edit_file"})
|
||||
|
||||
|
||||
def test_native_achat_follows_per_call_edit_tools(make, monkeypatch):
|
||||
import asyncio
|
||||
agent = make(llm="gpt-4o-mini")
|
||||
system, tool_names = _system_sent(agent, monkeypatch, lambda a: asyncio.run(a.achat("hi", tools=[EDIT_SCHEMA])))
|
||||
assert OPENAI_HINT in system
|
||||
assert "edit_file" in tool_names
|
||||
|
||||
|
||||
def test_system_prompt_attribute_states_instructions_once(make):
|
||||
agent = make(llm="gpt-4o-mini", instructions="Answer in French.")
|
||||
assert agent.system_prompt.count("Answer in French.") == 1
|
||||
|
||||
|
||||
def test_cached_prompt_follows_a_backstory_change(make):
|
||||
agent = make(llm="gpt-4o-mini", instructions="Answer in French.")
|
||||
assert "Your Goal:" not in agent._build_system_prompt()
|
||||
agent.backstory = "A translator"
|
||||
prompt = agent._build_system_prompt()
|
||||
assert "A translator" in prompt
|
||||
assert "Your Goal: Answer in French." in prompt
|
||||
+5
-1
@@ -32,6 +32,7 @@ def _agent_with_history(messages: int, words: int):
|
||||
|
||||
def fake_stream(**kwargs):
|
||||
sent["history"] = list(kwargs.get("chat_history") or [])
|
||||
sent["prompt"] = kwargs.get("prompt")
|
||||
yield "ok"
|
||||
|
||||
agent._using_custom_llm = True
|
||||
@@ -60,7 +61,10 @@ class TestCustomLlmStreamCompaction:
|
||||
"""Compaction must not fire when it is not needed."""
|
||||
agent, sent = _agent_with_history(6, 5)
|
||||
history = _drain(agent, sent)
|
||||
assert len(history) == 7, "a short history was compacted unnecessarily"
|
||||
# The six earlier turns go as history; this turn goes once, as the prompt.
|
||||
assert len(history) == 6, "a short history was compacted unnecessarily"
|
||||
assert sent["prompt"] == "hello"
|
||||
assert all(m["content"] != "hello" for m in history)
|
||||
|
||||
def test_the_context_manager_is_actually_configured_here(self):
|
||||
"""Guards the premise: tools present auto-enable context management."""
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
"""reasoning_effort must reach the provider on the Responses API path too.
|
||||
|
||||
OpenAI reasoning models (o1/o3/o4, gpt-5 on the native client) are sent to the
|
||||
Responses API by default. That API takes the level as
|
||||
``reasoning: {"effort": ...}``; neither Responses request builder applied the
|
||||
unified ``reasoning_effort`` (Issue #4452), so ``Agent(llm="o4-mini",
|
||||
reasoning_effort="high")`` produced a request without any effort on the
|
||||
default, non-streaming path. These tests assert on what the transport receives.
|
||||
"""
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("OPENAI_API_KEY", "sk-test-not-real")
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
from praisonaiagents.llm.openai_client import OpenAIClient
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolate_responses_endpoint(monkeypatch):
|
||||
"""The Responses path is only taken without a custom endpoint."""
|
||||
monkeypatch.delenv('OPENAI_API_BASE', raising=False)
|
||||
monkeypatch.delenv('OPENAI_BASE_URL', raising=False)
|
||||
|
||||
|
||||
def _response(text='ok'):
|
||||
return {'status': 'completed', 'output': [
|
||||
{'type': 'message', 'content': [{'type': 'output_text', 'text': text}]}]}
|
||||
|
||||
|
||||
MESSAGES = [{'role': 'user', 'content': 'question'}]
|
||||
|
||||
|
||||
class TestLLMResponsesParams:
|
||||
|
||||
@pytest.mark.parametrize('model', ['o4-mini', 'openai/o3', 'azure/o1'])
|
||||
def test_instance_effort_becomes_reasoning_effort_object(self, model):
|
||||
params = LLM(model=model, reasoning_effort='high')._build_responses_params(MESSAGES)
|
||||
assert params['reasoning'] == {'effort': 'high'}
|
||||
assert 'reasoning_effort' not in params
|
||||
|
||||
def test_per_call_effort_overrides_instance(self):
|
||||
llm = LLM(model='o4-mini', reasoning_effort='low')
|
||||
params = llm._build_responses_params(MESSAGES, reasoning_effort='high')
|
||||
assert params['reasoning'] == {'effort': 'high'}
|
||||
assert 'reasoning_effort' not in params
|
||||
|
||||
@pytest.mark.parametrize('effort', [None, 'off'])
|
||||
def test_unset_or_off_sends_nothing(self, effort):
|
||||
params = LLM(model='o4-mini', reasoning_effort=effort)._build_responses_params(MESSAGES)
|
||||
assert 'reasoning' not in params
|
||||
|
||||
def test_model_without_reasoning_control_is_untouched(self):
|
||||
params = LLM(model='gpt-4o-mini', reasoning_effort='high')._build_responses_params(MESSAGES)
|
||||
assert 'reasoning' not in params
|
||||
assert 'reasoning_effort' not in params
|
||||
|
||||
def test_explicit_reasoning_object_is_kept(self):
|
||||
llm = LLM(model='o4-mini', reasoning_effort='high')
|
||||
explicit = {'effort': 'low', 'summary': 'auto'}
|
||||
params = llm._build_responses_params(MESSAGES, reasoning=explicit)
|
||||
assert params['reasoning'] == explicit
|
||||
|
||||
def test_public_get_response_sends_effort(self):
|
||||
llm = LLM(model='o4-mini', reasoning_effort='high')
|
||||
requests = []
|
||||
|
||||
def respond(**kwargs):
|
||||
requests.append(kwargs)
|
||||
return _response()
|
||||
|
||||
llm._call_responses_api = respond
|
||||
assert llm.get_response('question', stream=False, verbose=False) == 'ok'
|
||||
assert requests[0]['reasoning'] == {'effort': 'high'}
|
||||
|
||||
|
||||
class TestOpenAIClientResponsesInput:
|
||||
|
||||
def test_reasoning_effort_kwarg_becomes_reasoning_object(self):
|
||||
client = OpenAIClient.__new__(OpenAIClient)
|
||||
params = client._build_responses_input(MESSAGES, 'o4-mini', None, None, reasoning_effort='high')
|
||||
assert params['reasoning'] == {'effort': 'high'}
|
||||
assert 'reasoning_effort' not in params
|
||||
|
||||
def test_no_effort_no_reasoning_object(self):
|
||||
client = OpenAIClient.__new__(OpenAIClient)
|
||||
params = client._build_responses_input(MESSAGES, 'o4-mini', None, None)
|
||||
assert 'reasoning' not in params
|
||||
|
||||
|
||||
def _agent_request(model, **agent_kwargs):
|
||||
"""Drive a real Agent on the native OpenAI client, faking only responses.create."""
|
||||
requests = []
|
||||
|
||||
def respond(**kwargs):
|
||||
requests.append(kwargs)
|
||||
return _response()
|
||||
|
||||
async def arespond(**kwargs):
|
||||
return respond(**kwargs)
|
||||
|
||||
client = OpenAIClient(api_key='sk-test-not-real')
|
||||
client._sync_client = SimpleNamespace(responses=SimpleNamespace(create=respond))
|
||||
client._async_client = SimpleNamespace(responses=SimpleNamespace(create=arespond))
|
||||
agent = Agent(name='t', instructions='x', llm=model, **agent_kwargs)
|
||||
agent._Agent__openai_client = client
|
||||
assert agent.start('hi') == 'ok'
|
||||
assert requests, 'no Responses API request was made'
|
||||
return requests[-1]
|
||||
|
||||
|
||||
class TestAgentNativeOpenAIPath:
|
||||
|
||||
def test_reasoning_model_gets_effort(self):
|
||||
assert _agent_request('o4-mini', reasoning_effort='high')['reasoning'] == {'effort': 'high'}
|
||||
|
||||
def test_unset_effort_sends_nothing(self):
|
||||
assert 'reasoning' not in _agent_request('o4-mini')
|
||||
|
||||
def test_model_without_reasoning_control_is_untouched(self):
|
||||
assert 'reasoning' not in _agent_request('gpt-4o-mini', reasoning_effort='high')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Automatic promotion must honor the same capacity as explicit long-term writes."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from praisonaiagents.memory.file_memory import FileMemory
|
||||
|
||||
|
||||
@pytest.mark.parametrize("limit", [1, 2, 3])
|
||||
def test_promotion_enforces_persisted_long_term_limit(tmp_path, limit):
|
||||
memory = FileMemory(
|
||||
base_path=str(tmp_path),
|
||||
config={"short_term_limit": 2, "long_term_limit": limit},
|
||||
)
|
||||
for index in range(6):
|
||||
memory.add_short_term(f"Important fact {index}", importance=0.8)
|
||||
assert len(memory.get_long_term()) == limit
|
||||
assert len(FileMemory(base_path=str(tmp_path)).get_long_term()) == limit
|
||||
assert len(memory.get_short_term()) == 2
|
||||
|
||||
|
||||
def test_promotion_keeps_highest_importance_and_newest_ties(tmp_path):
|
||||
memory = FileMemory(
|
||||
base_path=str(tmp_path),
|
||||
config={"short_term_limit": 3, "long_term_limit": 2},
|
||||
)
|
||||
memory.add_long_term("Keep strongest", importance=1.0)
|
||||
memory.add_long_term("Older tie", importance=0.8)
|
||||
memory.add_short_term("Lower priority", importance=0.7)
|
||||
memory.add_short_term("Newer tie", importance=0.8, metadata={"source": "fixture"})
|
||||
memory.add_short_term("Below promotion threshold", importance=0.2)
|
||||
for path, timestamp in [(memory.long_term_file, 100), (memory.short_term_file, 200)]:
|
||||
records = json.loads(path.read_text())
|
||||
for record in records:
|
||||
record["created_at"] = timestamp
|
||||
path.write_text(json.dumps(records))
|
||||
promoted_before = next(
|
||||
item for item in json.loads(memory.short_term_file.read_text())
|
||||
if item["content"] == "Newer tie"
|
||||
)
|
||||
memory.add_short_term("Trigger overflow", importance=0.2)
|
||||
contents = {item.content for item in memory.get_long_term()}
|
||||
assert contents == {"Keep strongest", "Newer tie"}
|
||||
promoted_after = next(item for item in memory.get_long_term() if item.content == "Newer tie")
|
||||
assert promoted_after.to_dict() == promoted_before
|
||||
assert {item.content for item in memory.get_short_term()} == {
|
||||
"Newer tie", "Below promotion threshold", "Trigger overflow",
|
||||
}
|
||||
assert {item.content for item in FileMemory(base_path=str(tmp_path)).get_long_term()} == contents
|
||||
|
||||
|
||||
def test_disabled_promotion_preserves_existing_long_term(tmp_path):
|
||||
memory = FileMemory(
|
||||
base_path=str(tmp_path),
|
||||
config={"short_term_limit": 1, "long_term_limit": 1, "auto_promote": False},
|
||||
)
|
||||
memory.add_long_term("Existing fact")
|
||||
for index in range(3):
|
||||
memory.add_short_term(f"Unpromoted {index}", importance=1.0)
|
||||
assert [item.content for item in memory.get_long_term()] == ["Existing fact"]
|
||||
|
||||
|
||||
def test_promotion_zero_capacity_matches_explicit_long_term(tmp_path):
|
||||
memory = FileMemory(
|
||||
base_path=str(tmp_path),
|
||||
config={"short_term_limit": 1, "long_term_limit": 0},
|
||||
)
|
||||
memory.add_long_term("Explicit fact", importance=1.0)
|
||||
assert memory.get_long_term() == []
|
||||
memory.add_short_term("Promoted one", importance=1.0)
|
||||
memory.add_short_term("Promoted two", importance=1.0)
|
||||
assert memory.get_long_term() == []
|
||||
assert FileMemory(base_path=str(tmp_path)).get_long_term() == []
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Tests for the session manifest (Issue #6248)."""
|
||||
|
||||
from praisonaiagents.session import FileSessionManifest, SessionManifestProtocol
|
||||
|
||||
|
||||
def test_satisfies_protocol(tmp_path):
|
||||
assert isinstance(FileSessionManifest(str(tmp_path)), SessionManifestProtocol)
|
||||
|
||||
|
||||
def test_missing_session_has_no_entry(tmp_path):
|
||||
m = FileSessionManifest(str(tmp_path))
|
||||
assert m.get("missing") is None and m.interrupted("missing") is None
|
||||
|
||||
|
||||
def test_legacy_entry_with_unknown_keys_still_loads(tmp_path):
|
||||
(tmp_path / "s1.json").write_text(
|
||||
'{"session_id": "s1", "refs": {"plan": "p"}, "run_in_flight": true}',
|
||||
encoding="utf-8",
|
||||
)
|
||||
assert FileSessionManifest(str(tmp_path)).interrupted("s1") is not None
|
||||
|
||||
|
||||
def test_completed_run_is_not_interrupted(tmp_path):
|
||||
m = FileSessionManifest(str(tmp_path))
|
||||
m.begin_run("s1", "r1", run_input="task")
|
||||
assert m.interrupted("s1") is not None
|
||||
m.end_run("s1", "completed")
|
||||
assert m.interrupted("s1") is None
|
||||
assert m.get("s1").run_status == "completed"
|
||||
|
||||
|
||||
def test_interrupted_run_survives_new_process(tmp_path):
|
||||
FileSessionManifest(str(tmp_path)).begin_run("s1", "r1", run_input="task")
|
||||
# A new process (new instance) sees the surviving write-ahead marker.
|
||||
stale = FileSessionManifest(str(tmp_path)).interrupted("s1")
|
||||
assert stale is not None and stale.run_id == "r1" and stale.run_digest
|
||||
|
||||
|
||||
def test_attempt_budget_is_bounded(tmp_path):
|
||||
m = FileSessionManifest(str(tmp_path))
|
||||
attempts = [m.begin_run("s1", f"r{i}", run_input="task", max_attempts=2).attempt
|
||||
for i in range(4)]
|
||||
# Retries of the same interrupted task count up, then the budget resets.
|
||||
assert attempts == [1, 2, 1, 2]
|
||||
|
||||
|
||||
def test_new_task_after_interruption_is_a_fresh_attempt(tmp_path):
|
||||
m = FileSessionManifest(str(tmp_path))
|
||||
m.begin_run("s1", "r1", run_input="task a")
|
||||
assert m.begin_run("s1", "r2", run_input="task b").attempt == 1
|
||||
|
||||
|
||||
def test_unsafe_session_id_stays_in_directory(tmp_path):
|
||||
m = FileSessionManifest(str(tmp_path))
|
||||
m.begin_run("../../etc/x", "r1")
|
||||
assert all(p.parent == tmp_path for p in tmp_path.iterdir())
|
||||
assert m.get("../../etc/x").run_id == "r1"
|
||||
|
||||
|
||||
def test_run_input_is_never_written_in_plaintext(tmp_path):
|
||||
secret = "deploy with token sk-SECRET-12345"
|
||||
m = FileSessionManifest(str(tmp_path))
|
||||
m.begin_run("s1", "r1", run_input=secret)
|
||||
for f in tmp_path.iterdir():
|
||||
assert "sk-SECRET-12345" not in f.read_text(encoding="utf-8")
|
||||
# The digest still recognises a retry of the same task.
|
||||
assert m.begin_run("s1", "r2", run_input=secret).attempt == 2
|
||||
|
||||
|
||||
def test_absolute_and_backslash_ids_stay_in_directory(tmp_path):
|
||||
m = FileSessionManifest(str(tmp_path))
|
||||
for sid in ("/etc/passwd", "..\\..\\x", "..", "C:\\x"):
|
||||
m.begin_run(sid, sid)
|
||||
assert m.get(sid).run_id == sid
|
||||
assert all(p.parent == tmp_path for p in tmp_path.iterdir())
|
||||
|
||||
|
||||
def test_ids_differing_only_in_punctuation_stay_separate(tmp_path):
|
||||
m = FileSessionManifest(str(tmp_path))
|
||||
m.begin_run("a:b", "colon")
|
||||
m.begin_run("a_b", "underscore")
|
||||
assert m.get("a:b").run_id == "colon"
|
||||
assert m.get("a_b").run_id == "underscore"
|
||||
|
||||
|
||||
def test_end_run_only_releases_its_own_marker(tmp_path):
|
||||
m = FileSessionManifest(str(tmp_path))
|
||||
m.begin_run("s1", "r1", run_input="first")
|
||||
m.begin_run("s1", "r2", run_input="second") # second terminal takes over
|
||||
assert m.end_run("s1", "completed", run_id="r1") is None
|
||||
stale = m.interrupted("s1")
|
||||
assert stale is not None and stale.run_id == "r2"
|
||||
m.end_run("s1", "completed", run_id="r2")
|
||||
assert m.interrupted("s1") is None
|
||||
@@ -34,11 +34,11 @@ def _restore_approval_state():
|
||||
from praisonaiagents.approval import get_approval_registry, clear_approval_context
|
||||
|
||||
registry = get_approval_registry()
|
||||
previous = registry.get_backend()
|
||||
previous = registry._global_backend # not get_backend(): its ConsoleBackend fallback would be installed globally on restore
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
registry.set_backend(previous)
|
||||
registry._global_backend = previous
|
||||
clear_approval_context()
|
||||
|
||||
|
||||
@@ -154,7 +154,7 @@ def test_code_mode_changes_what_the_tool_can_do():
|
||||
|
||||
_probe_tool_name = "code_mode_probe"
|
||||
get_registry().register(code_mode_probe)
|
||||
previous = get_approval_registry().get_backend()
|
||||
previous = get_approval_registry()._global_backend # not get_backend(): its ConsoleBackend fallback would be installed globally on restore
|
||||
get_approval_registry().set_backend(AutoApproveBackend())
|
||||
try:
|
||||
safe = Agent(
|
||||
@@ -182,7 +182,7 @@ def test_code_mode_changes_what_the_tool_can_do():
|
||||
safe_result = safe_exec(code)
|
||||
unsafe_result = unsafe_exec(code)
|
||||
finally:
|
||||
get_approval_registry().set_backend(previous)
|
||||
get_approval_registry()._global_backend = previous
|
||||
clear_approval_context()
|
||||
# The tool registry is process-global. Leaving these probe tools in it
|
||||
# changed what later tests saw when they enumerated or dispatched
|
||||
@@ -221,7 +221,7 @@ def test_unsafe_code_mode_cannot_reach_tools_the_agent_was_not_granted():
|
||||
|
||||
_probe_tool_name = "ungranted_global_tool"
|
||||
get_registry().register(ungranted_global_tool)
|
||||
previous = get_approval_registry().get_backend()
|
||||
previous = get_approval_registry()._global_backend # not get_backend(): its ConsoleBackend fallback would be installed globally on restore
|
||||
get_approval_registry().set_backend(AutoApproveBackend())
|
||||
try:
|
||||
agent = Agent(
|
||||
@@ -240,7 +240,7 @@ def test_unsafe_code_mode_cannot_reach_tools_the_agent_was_not_granted():
|
||||
)
|
||||
result = execute("print(ungranted_global_tool())")
|
||||
finally:
|
||||
get_approval_registry().set_backend(previous)
|
||||
get_approval_registry()._global_backend = previous
|
||||
clear_approval_context()
|
||||
# The tool registry is process-global. Leaving these probe tools in it
|
||||
# changed what later tests saw when they enumerated or dispatched
|
||||
@@ -279,13 +279,13 @@ def test_code_execution_runs_a_trailing_expression_exactly_once():
|
||||
from praisonaiagents.approval import clear_approval_context
|
||||
|
||||
registry = get_approval_registry()
|
||||
previous = registry.get_backend()
|
||||
previous = registry._global_backend # not get_backend(): its ConsoleBackend fallback would be installed globally on restore
|
||||
registry.set_backend(AutoApproveBackend())
|
||||
try:
|
||||
assert execute_code("print(7)")["stdout"] == "7\n"
|
||||
assert execute_code_with_tools("print(7)")["stdout"] == "7\n"
|
||||
finally:
|
||||
registry.set_backend(previous)
|
||||
registry._global_backend = previous
|
||||
clear_approval_context()
|
||||
|
||||
|
||||
|
||||
@@ -79,7 +79,7 @@ def test_no_handrolled_async_bridges_outside_helper():
|
||||
for path in root.rglob("*.py"):
|
||||
if path in allowed:
|
||||
continue
|
||||
tree = ast.parse(path.read_text(), filename=str(path))
|
||||
tree = ast.parse(path.read_bytes(), filename=str(path))
|
||||
for func in ast.walk(tree):
|
||||
if not isinstance(func, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
continue
|
||||
@@ -93,6 +93,15 @@ def test_no_handrolled_async_bridges_outside_helper():
|
||||
)
|
||||
|
||||
|
||||
def test_scan_accepts_python_source_encoding_cookie(tmp_path, monkeypatch):
|
||||
"""Scan the same Latin-1 source that Python's parser accepts."""
|
||||
(tmp_path / "example.py").write_bytes(
|
||||
b"# coding: latin-1; caf\xe9\nmessage = 'caf\xe9'\n"
|
||||
)
|
||||
monkeypatch.setitem(globals(), "_package_root", lambda: tmp_path)
|
||||
test_no_handrolled_async_bridges_outside_helper()
|
||||
|
||||
|
||||
def test_tool_timeout_uses_daemon_worker_and_does_not_leak():
|
||||
before = {t.ident for t in threading.enumerate()}
|
||||
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Counting rules of benchmarks/execution_benchmark.py, with fake runners.
|
||||
|
||||
No framework is imported and nothing leaves the machine: each runner is a
|
||||
``build()`` that returns the call to time, as the script's real runners do.
|
||||
"""
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
SCRIPT = Path(__file__).resolve().parents[2] / "benchmarks" / "execution_benchmark.py"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def bench():
|
||||
spec = importlib.util.spec_from_file_location("execution_benchmark", SCRIPT)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def runner(answers, log, name):
|
||||
"""A build() whose calls return ``answers`` in order (an Exception is raised)."""
|
||||
answers = iter(answers)
|
||||
|
||||
def build():
|
||||
log.append(("build", name))
|
||||
|
||||
def call():
|
||||
log.append(("call", name))
|
||||
answer = next(answers)
|
||||
if isinstance(answer, Exception):
|
||||
raise answer
|
||||
return answer
|
||||
return call
|
||||
return build
|
||||
|
||||
|
||||
def test_warm_up_calls_are_not_timed_and_rounds_interleave(bench):
|
||||
log = []
|
||||
runners = {"A": runner(["w", "a1", "a2"], log, "A"), "B": runner(["w", "b1", "b2"], log, "B")}
|
||||
results = bench.run_benchmark(runners, iterations=2, warmup=1)
|
||||
assert [len(r["times"]) for r in results.values()] == [2, 2]
|
||||
calls = [name for kind, name in log if kind == "call"]
|
||||
assert calls == ["A", "B", "A", "B", "B", "A"] # warm-ups, round 1, round 2 rotated
|
||||
|
||||
|
||||
def test_errors_and_empty_answers_are_failures_not_timings(bench):
|
||||
log = []
|
||||
runners = {"A": runner(["w", RuntimeError("401"), None, "None", " ", "4"], log, "A")}
|
||||
results = bench.run_benchmark(runners, iterations=5, warmup=1)
|
||||
assert len(results["A"]["times"]) == 1
|
||||
assert len(results["A"]["failures"]) == 4
|
||||
assert "RuntimeError: 401" in results["A"]["failures"][0]
|
||||
|
||||
|
||||
def test_a_failed_warm_up_does_not_drop_the_framework(bench):
|
||||
log = []
|
||||
runners = {"A": runner([RuntimeError("timeout"), "a1", "a2"], log, "A")}
|
||||
results = bench.run_benchmark(runners, iterations=2, warmup=1)
|
||||
assert len(results["A"]["times"]) == 2
|
||||
|
||||
|
||||
def test_summary_ranks_by_median_and_skips_frameworks_without_a_success(bench):
|
||||
results = {
|
||||
"slow": {"times": [1.0, 1.1, 9.0], "failures": []},
|
||||
"fast": {"times": [0.5, 0.6, 0.7], "failures": ["x"]},
|
||||
"broken": {"times": [], "failures": ["401"]},
|
||||
}
|
||||
rows = bench.summarise(results)
|
||||
assert [r[0] for r in rows] == ["fast", "slow"]
|
||||
name, median, mean, ok, failed = rows[1]
|
||||
assert (median, ok, failed) == (1.1, 3, 0)
|
||||
assert mean == pytest.approx(11.1 / 3)
|
||||
|
||||
|
||||
def test_only_the_call_is_timed(bench, monkeypatch):
|
||||
clock = iter([10.0, 10.5]) # start / end of the one timed call; build() reads no clock
|
||||
|
||||
def build():
|
||||
return lambda: "4"
|
||||
|
||||
monkeypatch.setattr(bench.time, "perf_counter", lambda: next(clock))
|
||||
results = bench.run_benchmark({"A": build}, iterations=1, warmup=0)
|
||||
assert results["A"]["times"] == [0.5]
|
||||
@@ -328,7 +328,7 @@ class TestMultiChannelGatewayConfig:
|
||||
from praisonaiagents.gateway.config import ChannelRouteConfig, MultiChannelGatewayConfig
|
||||
module = sys.modules["praisonaiagents.gateway.config"]
|
||||
source_file = module.__file__
|
||||
with open(source_file, "r") as f:
|
||||
with open(source_file, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
for dep in ["chromadb", "fastapi", "uvicorn", "pydantic", "litellm"]:
|
||||
assert f"import {dep}" not in content, f"Found heavy dep '{dep}' in config.py"
|
||||
@@ -376,7 +376,7 @@ class TestChatCommandInfo:
|
||||
from praisonaiagents.bots.protocols import BotProtocol
|
||||
module = sys.modules["praisonaiagents.bots.protocols"]
|
||||
source_file = module.__file__
|
||||
with open(source_file, "r") as f:
|
||||
with open(source_file, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
for dep in ["chromadb", "fastapi", "uvicorn", "litellm"]:
|
||||
assert f"import {dep}" not in content, f"Found heavy dep '{dep}' in protocols.py"
|
||||
|
||||
@@ -9,7 +9,7 @@ from praisonaiagents import ImageAgent
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_kwargs,expected", [
|
||||
({}, {"generate": "dall-e-3", "edit": "dall-e-2", "variation": "dall-e-2"}),
|
||||
({}, {"generate": "gpt-image-2.5-flare", "edit": "gpt-image-2.5-flare", "variation": "dall-e-2"}),
|
||||
({"llm": "openai/gpt-image-1"}, dict.fromkeys(("generate", "edit", "variation"), "openai/gpt-image-1")),
|
||||
({"model": "openai/gpt-image-1"}, dict.fromkeys(("generate", "edit", "variation"), "openai/gpt-image-1")),
|
||||
])
|
||||
@@ -37,3 +37,27 @@ def test_image_model_resolution(monkeypatch, tmp_path, model_kwargs, expected):
|
||||
agent.variation(str(image))
|
||||
|
||||
assert {name: kwargs["model"] for name, kwargs in calls.items()} == expected
|
||||
|
||||
|
||||
IMAGES_26_PARAMS = {"quality": "low", "frame_count": 4, "frame_duration_ms": 100,
|
||||
"complexity": "low", "output_format": "gif"}
|
||||
|
||||
|
||||
def test_generate_forwards_images_26_params(monkeypatch):
|
||||
calls = {}
|
||||
agent = ImageAgent(llm="gpt-image-2.6-flare", verbose=False)
|
||||
monkeypatch.setattr(agent, "_litellm", lambda **kw: calls.update(kw) or {"data": []})
|
||||
monkeypatch.setitem(sys.modules, "litellm", types.SimpleNamespace(
|
||||
get_llm_provider=lambda model: (model, "openai", None, None)))
|
||||
agent.generate_image("x", **IMAGES_26_PARAMS)
|
||||
assert IMAGES_26_PARAMS.items() <= calls.items()
|
||||
|
||||
|
||||
def test_edit_rejects_markup_instead_of_dropping_it(monkeypatch):
|
||||
called = []
|
||||
agent = ImageAgent(llm="gpt-image-2.5-flare", verbose=False)
|
||||
monkeypatch.setitem(sys.modules, "litellm", types.SimpleNamespace(
|
||||
image_edit=lambda **kw: called.append(kw)))
|
||||
with pytest.raises(ValueError, match="markup"):
|
||||
agent.edit("photo.png", "make the circled cup blue", markup="markup.png")
|
||||
assert not called
|
||||
@@ -0,0 +1,158 @@
|
||||
"""ImageAgent results expose .save() while keeping the raw response API (#6260)."""
|
||||
|
||||
import base64
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import ImageAgent
|
||||
|
||||
GIF = b"GIF89a\x01\x00\x01\x00fake"
|
||||
SVG = b'<svg xmlns="http://www.w3.org/2000/svg"></svg>'
|
||||
|
||||
|
||||
def _response(*payloads):
|
||||
return types.SimpleNamespace(
|
||||
data=[types.SimpleNamespace(b64_json=base64.b64encode(p).decode(), url=None) for p in payloads],
|
||||
usage={"total_tokens": 5},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent(monkeypatch):
|
||||
agent = ImageAgent(llm="openai/gpt-image-2.6-flare", verbose=False)
|
||||
sent = {}
|
||||
|
||||
def fake(name, payloads):
|
||||
def call(**kwargs):
|
||||
sent[name] = kwargs
|
||||
return _response(*payloads)
|
||||
return call
|
||||
|
||||
monkeypatch.setattr(agent, "_litellm", fake("generate", [GIF]))
|
||||
monkeypatch.setitem(sys.modules, "litellm", types.SimpleNamespace(
|
||||
image_edit=fake("edit", [SVG, GIF]),
|
||||
image_variation=fake("variation", [GIF]),
|
||||
get_llm_provider=lambda model: (model, "openai", None, None),
|
||||
))
|
||||
agent.sent = sent
|
||||
return agent
|
||||
|
||||
|
||||
def test_generate_save_writes_decoded_bytes(agent, tmp_path):
|
||||
out = agent.generate("A cat waving", output_format="gif", frame_count=8).save(str(tmp_path / "cat.gif"))
|
||||
assert open(out, "rb").read() == GIF
|
||||
assert agent.sent["generate"]["output_format"] == "gif"
|
||||
assert agent.sent["generate"]["frame_count"] == 8
|
||||
|
||||
|
||||
def test_edit_save_index_and_raw_access(agent, tmp_path):
|
||||
src = tmp_path / "photo.png"
|
||||
src.write_bytes(b"png")
|
||||
result = agent.edit(str(src), "make the cup blue", n=2)
|
||||
assert agent.sent["edit"]["n"] == 2
|
||||
assert len(result.images) == 2
|
||||
assert result.usage == {"total_tokens": 5}
|
||||
assert base64.b64decode(result.data[1].b64_json) == GIF
|
||||
result.save(str(tmp_path / "a.svg"))
|
||||
assert (tmp_path / "a.svg").read_bytes() == SVG
|
||||
|
||||
|
||||
def test_variation_and_dict_response(agent, tmp_path, monkeypatch):
|
||||
src = tmp_path / "photo.png"
|
||||
src.write_bytes(b"png")
|
||||
assert agent.variation(str(src)).save(str(tmp_path / "v.gif"))
|
||||
monkeypatch.setattr(agent, "_litellm", lambda **kw: {"data": [{"b64_json": base64.b64encode(SVG).decode()}]})
|
||||
result = agent.generate("logo")
|
||||
assert result["data"][0]["b64_json"]
|
||||
result.save(str(tmp_path / "logo.svg"))
|
||||
assert (tmp_path / "logo.svg").read_bytes() == SVG
|
||||
|
||||
|
||||
def test_save_without_image_data_raises(agent, tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(agent, "_litellm", lambda **kw: {"data": []})
|
||||
with pytest.raises(ValueError):
|
||||
agent.generate("nothing").save(str(tmp_path / "x.png"))
|
||||
|
||||
|
||||
def test_save_nonzero_index(agent, tmp_path):
|
||||
src = tmp_path / "photo.png"
|
||||
src.write_bytes(b"png")
|
||||
result = agent.edit(str(src), "two takes", n=2)
|
||||
assert result.save(str(tmp_path / "b.gif"), index=1) == str(tmp_path / "b.gif")
|
||||
assert (tmp_path / "b.gif").read_bytes() == GIF
|
||||
|
||||
|
||||
class _FakeHTTP:
|
||||
def __init__(self, body):
|
||||
self.body = body
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc):
|
||||
return False
|
||||
|
||||
def read(self, n=-1):
|
||||
return self.body if n < 0 else self.body[:n]
|
||||
|
||||
|
||||
def _url_response(url):
|
||||
return {"data": [{"b64_json": None, "url": url}]}
|
||||
|
||||
|
||||
def test_save_downloads_url(agent, tmp_path, monkeypatch):
|
||||
import urllib.request
|
||||
|
||||
seen = {}
|
||||
|
||||
def fake_urlopen(url, timeout=None):
|
||||
seen.update(url=url, timeout=timeout)
|
||||
return _FakeHTTP(GIF)
|
||||
|
||||
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
|
||||
monkeypatch.setattr(agent, "_litellm", lambda **kw: _url_response("https://example.com/i.gif"))
|
||||
agent.generate("cat").save(str(tmp_path / "u.gif"))
|
||||
assert (tmp_path / "u.gif").read_bytes() == GIF
|
||||
assert seen == {"url": "https://example.com/i.gif", "timeout": 30}
|
||||
|
||||
|
||||
def test_save_rejects_non_http_url(agent, tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(agent, "_litellm", lambda **kw: _url_response("file:///etc/passwd"))
|
||||
with pytest.raises(ValueError, match="http"):
|
||||
agent.generate("x").save(str(tmp_path / "x.png"))
|
||||
assert not (tmp_path / "x.png").exists()
|
||||
|
||||
|
||||
def test_save_rejects_oversized_download(agent, tmp_path, monkeypatch):
|
||||
import urllib.request
|
||||
from praisonaiagents.agent import _image_result
|
||||
|
||||
monkeypatch.setattr(_image_result, "_MAX_DOWNLOAD_BYTES", 4)
|
||||
monkeypatch.setattr(urllib.request, "urlopen", lambda url, timeout=None: _FakeHTTP(b"123456"))
|
||||
monkeypatch.setattr(agent, "_litellm", lambda **kw: _url_response("https://example.com/big.png"))
|
||||
with pytest.raises(ValueError, match="maximum"):
|
||||
agent.generate("x").save(str(tmp_path / "big.png"))
|
||||
assert not (tmp_path / "big.png").exists()
|
||||
|
||||
|
||||
def test_container_protocol_passes_through(agent):
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
raw = ImageResponse(data=[{"b64_json": base64.b64encode(GIF).decode()}])
|
||||
agent._litellm = lambda **kw: raw
|
||||
result = agent.generate("cat")
|
||||
assert "data" in result
|
||||
assert dict(result).keys() == dict(raw).keys()
|
||||
|
||||
|
||||
def test_workflow_output_saves_b64_image(agent, tmp_path):
|
||||
from praisonaiagents.workflows.workflows import AgentFlow
|
||||
|
||||
flow = object.__new__(AgentFlow)
|
||||
out = str(tmp_path / "wf.gif")
|
||||
variables = {"output": out}
|
||||
flow._execute_specialized_agent(agent, "ImageAgent", "a cat", None, variables, False)
|
||||
assert (tmp_path / "wf.gif").read_bytes() == GIF
|
||||
assert variables["_last_image_file"] == out
|
||||
@@ -0,0 +1,107 @@
|
||||
"""Regression tests for Issue #6238.
|
||||
|
||||
* Caller-supplied names must never escape the storage directory.
|
||||
* Sync entry points must work from inside a running event loop.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from praisonaiagents.planning.storage import PlanStorage
|
||||
from praisonaiagents.planning.todo import TodoList
|
||||
from praisonaiagents.utils.atomic_io import safe_join
|
||||
|
||||
|
||||
def test_safe_join_keeps_legacy_names_and_confines_others(tmp_path):
|
||||
base = str(tmp_path)
|
||||
assert safe_join(base, "current", ".json") == os.path.join(base, "current.json")
|
||||
for name in ("../../x", "/etc/passwd", "..", "", "a/b", "a\\b", ".hidden"):
|
||||
path = safe_join(base, name, ".json")
|
||||
assert os.path.dirname(path) == base, name
|
||||
# Distinct inputs never collide on one file.
|
||||
assert safe_join(base, "a/b") != safe_join(base, "a_b")
|
||||
|
||||
|
||||
def test_plan_storage_todo_name_cannot_traverse(tmp_path):
|
||||
storage = PlanStorage(base_path=str(tmp_path / "store"))
|
||||
path = storage.save_todo(TodoList(), name="../../escaped")
|
||||
assert os.path.dirname(path) == storage.todos_dir
|
||||
assert not (tmp_path / "escaped.json").exists()
|
||||
assert storage.load_todo("../../escaped") is not None
|
||||
assert storage.delete_todo("../../escaped") is True
|
||||
|
||||
|
||||
def test_generate_title_works_inside_running_loop(monkeypatch):
|
||||
from praisonaiagents.session import title
|
||||
|
||||
async def fake(*args, **kwargs):
|
||||
return "LLM Title"
|
||||
|
||||
monkeypatch.setattr(title, "generate_title_async", fake)
|
||||
|
||||
async def main():
|
||||
return title.generate_title("hello", "world")
|
||||
|
||||
assert asyncio.run(main()) == "LLM Title"
|
||||
|
||||
|
||||
def test_listed_todo_names_round_trip(tmp_path):
|
||||
storage = PlanStorage(base_path=str(tmp_path / "store"))
|
||||
storage.save_todo(TodoList(), name="my tasks")
|
||||
(listed,) = storage.list_todos()
|
||||
assert storage.load_todo(listed) is not None
|
||||
assert storage.load_todo("my tasks") is not None
|
||||
assert storage.delete_todo(listed) is True
|
||||
assert storage.list_todos() == []
|
||||
|
||||
|
||||
def test_legacy_todo_file_still_found(tmp_path):
|
||||
storage = PlanStorage(base_path=str(tmp_path / "store"))
|
||||
legacy = os.path.join(storage.todos_dir, "my tasks.json")
|
||||
with open(legacy, "w", encoding="utf-8") as f:
|
||||
f.write(TodoList().to_json())
|
||||
assert storage.load_todo("my tasks") is not None
|
||||
assert storage.save_todo(TodoList(), name="my tasks") == legacy
|
||||
assert storage.list_todos() == ["my tasks"]
|
||||
assert storage.delete_todo("my tasks") is True
|
||||
|
||||
|
||||
def test_custom_plan_filename_keeps_md_suffix(tmp_path):
|
||||
from praisonaiagents.planning.plan import Plan
|
||||
|
||||
storage = PlanStorage(base_path=str(tmp_path / "store"))
|
||||
plan = Plan(name="demo")
|
||||
path = storage.save_plan(plan, filename="my plan.md")
|
||||
assert path.endswith(".md") and os.path.dirname(path) == storage.plans_dir
|
||||
assert len(storage.list_plans()) == 1
|
||||
|
||||
|
||||
def test_encoded_looking_names_do_not_share_a_file(tmp_path):
|
||||
storage = PlanStorage(base_path=str(tmp_path / "store"))
|
||||
names = ["a/b", "a_b~c14cddc033f6", "a%2Fb", "a b", "a%20b", "~x", ".."]
|
||||
paths = {storage.save_todo(TodoList(name=n), name=n) for n in names}
|
||||
assert len(paths) == len(names)
|
||||
assert all(os.path.dirname(p) == storage.todos_dir for p in paths)
|
||||
assert sorted(storage.list_todos()) == sorted(names)
|
||||
for n in names:
|
||||
assert storage.load_todo(n).name == n
|
||||
assert storage.delete_todo("a/b") is True
|
||||
assert storage.load_todo("a%2Fb") is not None
|
||||
|
||||
|
||||
def test_long_unicode_names_are_bounded_and_round_trip(tmp_path):
|
||||
storage = PlanStorage(base_path=str(tmp_path / "store"))
|
||||
long_a, long_b = "我的任务" * 40, "我的任务" * 40 + "x"
|
||||
pa = storage.save_todo(TodoList(name=long_a), name=long_a)
|
||||
pb = storage.save_todo(TodoList(name=long_b), name=long_b)
|
||||
assert pa != pb
|
||||
for p in (pa, pb):
|
||||
assert os.path.dirname(p) == storage.todos_dir
|
||||
assert len(os.path.basename(p).encode("utf-8")) <= 200
|
||||
# Short non-ASCII names keep their plain (pre-#6238) filename.
|
||||
assert storage.save_todo(TodoList(), name="我的任务").endswith("我的任务.json")
|
||||
for listed in storage.list_todos():
|
||||
assert storage.load_todo(listed) is not None
|
||||
assert storage.delete_todo(listed) is True
|
||||
assert storage.list_todos() == []
|
||||
assert storage.load_todo(long_a) is None
|
||||
@@ -0,0 +1,62 @@
|
||||
"""Regression tests for issue #6257 (beginner harness demo)."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.llm.openai_client import OpenAIClient
|
||||
|
||||
|
||||
def test_guardrail_runs_when_agent_has_no_name():
|
||||
seen = []
|
||||
|
||||
def check(output):
|
||||
seen.append(output.agent)
|
||||
return True, output
|
||||
|
||||
agent = Agent(instructions="test", llm="gpt-4o-mini", guardrails=check)
|
||||
assert agent.name is None or isinstance(agent.name, str)
|
||||
|
||||
ok, result, error = agent._validate_with_guardrail("hello")
|
||||
|
||||
assert ok is True and error is None
|
||||
assert result == "hello"
|
||||
assert seen and isinstance(seen[0], str)
|
||||
|
||||
|
||||
def test_openai_timeout_env_var(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_TIMEOUT", "45")
|
||||
assert OpenAIClient(api_key="sk-test").timeout == 45.0
|
||||
|
||||
monkeypatch.setenv("OPENAI_TIMEOUT", "not-a-number")
|
||||
assert OpenAIClient(api_key="sk-test").timeout is None
|
||||
|
||||
monkeypatch.delenv("OPENAI_TIMEOUT")
|
||||
assert OpenAIClient(api_key="sk-test").timeout is None
|
||||
|
||||
|
||||
def test_openai_timeout_reaches_sync_client(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_TIMEOUT", "12")
|
||||
client = OpenAIClient(api_key="sk-test")
|
||||
monkeypatch.setattr(
|
||||
"praisonaiagents.llm.openai_client.check_model_request", lambda *a, **k: None
|
||||
)
|
||||
assert client.sync_client.timeout == 12.0
|
||||
|
||||
|
||||
def test_rules_notice_lists_only_active_rules_once(tmp_path, monkeypatch, caplog):
|
||||
import logging
|
||||
from praisonaiagents.agent import agent as agent_mod
|
||||
|
||||
rules_dir = tmp_path / ".praisonai" / "rules"
|
||||
rules_dir.mkdir(parents=True)
|
||||
(rules_dir / "always_on.md").write_text("---\nactivation: always\n---\nBe brief.\n")
|
||||
(rules_dir / "manual_only.md").write_text("---\nactivation: manual\n---\nOnly on @mention.\n")
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.setenv("HOME", str(tmp_path / "home"))
|
||||
monkeypatch.setattr(agent_mod, "_RULES_NOTICE_SEEN", set())
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
for _ in range(2):
|
||||
a = Agent(instructions="test", llm="gpt-4o-mini")
|
||||
a._init_rules_manager()
|
||||
|
||||
notices = [r.getMessage() for r in caplog.records if "Applied rules" in r.getMessage()]
|
||||
assert len(notices) == 1
|
||||
assert "always_on" in notices[0] and "manual_only" not in notices[0]
|
||||
@@ -64,6 +64,35 @@ def test_grep_glob_filters_by_pattern(sample_tree):
|
||||
assert "src" not in out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pattern", ["**/*.py", "src/**/*.py", "src/**/**/test_*.py"])
|
||||
def test_grep_recursive_glob_matches_zero_directories(sample_tree, pattern):
|
||||
from praisonaiagents.tools import grep
|
||||
|
||||
direct = sample_tree / "src" / "test_direct.py"
|
||||
direct.write_text("needle\n")
|
||||
nested = sample_tree / "src" / "nested"
|
||||
nested.mkdir()
|
||||
(nested / "test_nested.py").write_text("needle\n")
|
||||
with patch("shutil.which", return_value=None):
|
||||
out = grep("needle", glob=pattern)
|
||||
assert "src/test_direct.py:1:" in out.replace(os.sep, "/")
|
||||
assert "src/nested/test_nested.py:1:" in out.replace(os.sep, "/")
|
||||
if pattern == "**/*.py":
|
||||
assert "a.py:2:" in out
|
||||
else:
|
||||
assert "a.py:2:" not in out
|
||||
if "test_" in pattern:
|
||||
assert "src/d.py:1:" not in out.replace(os.sep, "/")
|
||||
|
||||
|
||||
def test_grep_recursive_glob_single_file(sample_tree):
|
||||
from praisonaiagents.tools import grep
|
||||
|
||||
with patch("shutil.which", return_value=None):
|
||||
out = grep("needle", path="a.py", glob="**/*.py")
|
||||
assert "a.py:2:" in out
|
||||
|
||||
|
||||
def test_grep_case_insensitive(sample_tree):
|
||||
from praisonaiagents.tools import grep
|
||||
with patch("shutil.which", return_value=None):
|
||||
|
||||
@@ -491,7 +491,7 @@ class TestTeamTasks:
|
||||
return ApprovalDecision(approved=True, reason="looks good", approver="ana")
|
||||
|
||||
registry = get_approval_registry()
|
||||
previous = registry.get_backend()
|
||||
previous = registry._global_backend # not get_backend(): its ConsoleBackend fallback would be installed globally on restore
|
||||
registry.set_backend(Reviewer())
|
||||
try:
|
||||
writer = Agent(name="Writer", instructions="write")
|
||||
@@ -500,7 +500,7 @@ class TestTeamTasks:
|
||||
out = TaskOutput(description="Draft", raw="api_key=sk-secret123456 text", agent="Writer")
|
||||
team._apply_human_review(task, 0, out)
|
||||
finally:
|
||||
registry.set_backend(previous)
|
||||
registry._global_backend = previous
|
||||
requested, decided = _of(recorder, "approval_requested"), _of(recorder, "approval_decided")
|
||||
assert len(requested) == len(decided) == 1
|
||||
assert requested[0].data["kind"] == "task_output_review"
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
"""Durable step cache: an interrupted AgentFlow resumes in a new process (#6249)."""
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.workflows.step_cache import FileStepCache
|
||||
from praisonaiagents.workflows.workflows import AgentFlow
|
||||
|
||||
|
||||
def _step(name, calls, output=None, fail=False):
|
||||
def handler(ctx):
|
||||
calls.append(name)
|
||||
if fail:
|
||||
raise RuntimeError("interrupted")
|
||||
return output if output is not None else f"{name}-out"
|
||||
handler.__name__ = name
|
||||
return handler
|
||||
|
||||
|
||||
def test_resume_skips_completed_steps_across_instances(tmp_path):
|
||||
path = tmp_path / "run.json"
|
||||
calls = []
|
||||
first = AgentFlow(
|
||||
steps=[_step("a", calls), _step("b", calls, fail=True)],
|
||||
cache=FileStepCache(str(path)),
|
||||
)
|
||||
try:
|
||||
first.run("in", verbose=False)
|
||||
except Exception:
|
||||
pass
|
||||
assert calls[0] == "a" and "b" in calls
|
||||
assert path.exists()
|
||||
|
||||
# A fresh flow and a fresh cache object: simulates a new process.
|
||||
calls.clear()
|
||||
second = AgentFlow(
|
||||
steps=[_step("a", calls), _step("b", calls), _step("c", calls)],
|
||||
cache=FileStepCache(str(path)),
|
||||
)
|
||||
result = second.run("in", verbose=False)
|
||||
assert calls == ["b", "c"]
|
||||
assert result["output"] == "c-out"
|
||||
|
||||
|
||||
def test_changed_input_does_not_serve_stale_results(tmp_path):
|
||||
path = str(tmp_path / "run.json")
|
||||
calls = []
|
||||
AgentFlow(steps=[_step("a", calls)], cache=FileStepCache(path)).run("x", verbose=False)
|
||||
AgentFlow(steps=[_step("a", calls)], cache=FileStepCache(path)).run("y", verbose=False)
|
||||
assert calls == ["a", "a"]
|
||||
|
||||
|
||||
def test_changed_upstream_output_reruns_downstream(tmp_path):
|
||||
path = str(tmp_path / "run.json")
|
||||
calls = []
|
||||
AgentFlow(
|
||||
steps=[_step("a", calls, output="v1"), _step("b", calls)], cache=FileStepCache(path)
|
||||
).run("x", verbose=False)
|
||||
calls.clear()
|
||||
# Step "a" was reconfigured under the SAME name (its closure value
|
||||
# changed), so it and everything after it must re-run.
|
||||
AgentFlow(
|
||||
steps=[_step("a", calls, output="v2"), _step("b", calls)], cache=FileStepCache(path)
|
||||
).run("x", verbose=False)
|
||||
assert calls == ["a", "b"]
|
||||
|
||||
|
||||
def test_unserialisable_value_is_rejected_not_stored(tmp_path):
|
||||
path = tmp_path / "run.json"
|
||||
cache = FileStepCache(str(path))
|
||||
with pytest.raises(TypeError):
|
||||
cache.set("k", {"output": object()})
|
||||
assert cache.get("k") is None and not path.exists()
|
||||
|
||||
|
||||
def test_get_returns_a_copy_and_clear_removes_file(tmp_path):
|
||||
path = tmp_path / "run.json"
|
||||
cache = FileStepCache(str(path))
|
||||
cache.set("k", {"output": "o", "variables": {"x": 1}})
|
||||
cache.get("k")["variables"]["x"] = 2
|
||||
assert FileStepCache(str(path)).get("k") == {"output": "o", "variables": {"x": 1}}
|
||||
assert json.loads(path.read_text())["k"]["output"] == "o"
|
||||
cache.clear()
|
||||
assert len(cache) == 0 and not path.exists()
|
||||
|
||||
|
||||
def test_edited_handler_with_same_name_reruns(tmp_path):
|
||||
path = str(tmp_path / "run.json")
|
||||
calls = []
|
||||
|
||||
def a(ctx):
|
||||
calls.append("v1")
|
||||
return "v1"
|
||||
|
||||
AgentFlow(steps=[a], cache=FileStepCache(path)).run("x", verbose=False)
|
||||
|
||||
def a(ctx): # noqa: F811 - same name, edited body
|
||||
calls.append("v2")
|
||||
return "v2"
|
||||
|
||||
result = AgentFlow(steps=[a], cache=FileStepCache(path)).run("x", verbose=False)
|
||||
assert calls == ["v1", "v2"] and result["output"] == "v2"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [{"variables": {"lookup": {1: "one"}}}, {"output": ("a", "b")}])
|
||||
def test_value_that_changes_shape_in_json_is_rejected(tmp_path, value):
|
||||
cache = FileStepCache(str(tmp_path / "run.json"))
|
||||
with pytest.raises(TypeError):
|
||||
cache.set("k", value)
|
||||
assert cache.get("k") is None
|
||||
|
||||
|
||||
def test_two_writers_on_one_path_keep_both_entries(tmp_path):
|
||||
path = str(tmp_path / "run.json")
|
||||
a, b = FileStepCache(path), FileStepCache(path)
|
||||
a.set("a", {"output": "1"})
|
||||
b.set("b", {"output": "2"})
|
||||
reopened = FileStepCache(path)
|
||||
assert reopened.get("a") == {"output": "1"} and reopened.get("b") == {"output": "2"}
|
||||
|
||||
|
||||
|
||||
def test_set_constants_are_encoded_independent_of_iteration_order():
|
||||
# A set literal compiles to a frozenset constant whose iteration order
|
||||
# depends on PYTHONHASHSEED; the key must not, or resume silently re-runs.
|
||||
from praisonaiagents.workflows.step_cache import _stable, make_step_key
|
||||
|
||||
words = ["apple", "banana", "cherry", "date", "elder"]
|
||||
expected = "{" + ",".join(sorted(repr(w) for w in words)) + "}"
|
||||
assert _stable(frozenset(words), set()) == expected
|
||||
assert _stable(frozenset(reversed(words)), set()) == expected
|
||||
|
||||
def lookup(ctx):
|
||||
return "fruit" if ctx.input in {"apple", "banana", "cherry", "date", "elder"} else "other"
|
||||
|
||||
consts = lookup.__code__.co_consts
|
||||
assert any(isinstance(c, frozenset) for c in consts)
|
||||
# Same constants, frozenset rebuilt: the key must not change.
|
||||
rebuilt = lookup.__code__.replace(
|
||||
co_consts=tuple(frozenset(sorted(c, reverse=True)) if isinstance(c, frozenset) else c for c in consts)
|
||||
)
|
||||
import types
|
||||
twin = types.FunctionType(rebuilt, lookup.__globals__, "lookup")
|
||||
assert make_step_key(lookup, None, "apple", {}) == make_step_key(twin, None, "apple", {})
|
||||
assert make_step_key(lookup, None, "x", {"t": {"b", "a"}}) == make_step_key(lookup, None, "x", {"t": {"a", "b"}})
|
||||
|
||||
def test_edited_helper_or_default_reruns(tmp_path):
|
||||
from praisonaiagents.workflows.step_cache import make_step_key
|
||||
|
||||
def helper_v1():
|
||||
return 1
|
||||
|
||||
def helper_v2():
|
||||
return 2
|
||||
|
||||
def make(default):
|
||||
def step(ctx, mode=default):
|
||||
return mode
|
||||
return step
|
||||
|
||||
assert make_step_key(make("fast"), None, "x", {}) != make_step_key(make("slow"), None, "x", {})
|
||||
assert make_step_key(make("fast"), None, "x", {}) == make_step_key(make("fast"), None, "x", {})
|
||||
|
||||
import types
|
||||
|
||||
def step(ctx):
|
||||
return helper() # noqa: F821 - resolved from the globals given below
|
||||
|
||||
def bind(helper):
|
||||
return types.FunctionType(step.__code__, {"helper": helper}, "step")
|
||||
|
||||
assert make_step_key(bind(helper_v1), None, "x", {}) != make_step_key(bind(helper_v2), None, "x", {})
|
||||
assert make_step_key(bind(helper_v1), None, "x", {}) == make_step_key(bind(helper_v1), None, "x", {})
|
||||
|
||||
|
||||
def test_agent_step_with_same_name_but_new_instructions_misses():
|
||||
from praisonaiagents.workflows.step_cache import make_step_key
|
||||
|
||||
class FakeAgent:
|
||||
def __init__(self, instructions, llm="gpt-4o-mini"):
|
||||
self.name, self.instructions, self.llm, self.tools = "writer", instructions, llm, []
|
||||
|
||||
k = lambda a: make_step_key(a, None, "x", {})
|
||||
assert k(FakeAgent("be brief")) == k(FakeAgent("be brief"))
|
||||
assert k(FakeAgent("be brief")) != k(FakeAgent("be verbose"))
|
||||
assert k(FakeAgent("be brief")) != k(FakeAgent("be brief", llm="gpt-4o"))
|
||||
|
||||
|
||||
def test_agent_wrapped_in_task_keys_on_agent_config():
|
||||
# AgentFlow normalises an Agent step to Task(agent=...) before keying it,
|
||||
# so the wrapped agent's instructions and model must reach the key.
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.workflows.step_cache import make_step_key
|
||||
|
||||
def key(instructions, llm="gpt-4o-mini"):
|
||||
agent = Agent(name="writer", instructions=instructions, llm=llm)
|
||||
task = AgentFlow(steps=[agent])._normalize_single_step(agent, 0)
|
||||
assert task.agent is agent
|
||||
return make_step_key(task, None, "x", {})
|
||||
|
||||
assert key("be brief") == key("be brief")
|
||||
assert key("be brief") != key("be verbose")
|
||||
assert key("be brief") != key("be brief", llm="gpt-4o")
|
||||
@@ -157,11 +157,12 @@ class TestCorrectnessOnAHit:
|
||||
|
||||
def test_legacy_cache_entry_restores_output_variable_and_status(self):
|
||||
cache = InMemoryStepCache()
|
||||
cache.set(make_step_key("upstream", None, "same", {"input": "same"}), {
|
||||
calls = []
|
||||
upstream = _counting_step("upstream", calls)
|
||||
cache.set(make_step_key(upstream, None, "same", {"input": "same"}), {
|
||||
"output": "cached", "variables": {"upstream_output": "cached"},
|
||||
})
|
||||
calls = []
|
||||
flow = AgentFlow(steps=[_counting_step("upstream", calls)], cache=cache)
|
||||
flow = AgentFlow(steps=[upstream], cache=cache)
|
||||
result = flow.run("same", verbose=False)
|
||||
assert calls == []
|
||||
assert result["output"] == "cached"
|
||||
|
||||
Generated
+1
-1
@@ -3162,7 +3162,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.21"
|
||||
version = "1.7.22"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.20"
|
||||
__version__ = "0.1.21"
|
||||
@@ -873,6 +873,15 @@ class HealthMonitorSchema(BaseModel):
|
||||
breaker_cooldown_s: float = Field(120.0, ge=0)
|
||||
|
||||
|
||||
class OperatorGrantSchema(BaseModel):
|
||||
"""Per-token operator grant (Issue #6253). Unknown keys are rejected so a
|
||||
typo (``agnts``) cannot silently drop an agent confinement."""
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
scopes: List[str] = Field(default_factory=list)
|
||||
agents: Optional[List[str]] = None
|
||||
|
||||
|
||||
class GatewayServerSchema(BaseModel):
|
||||
"""Typed schema for the ``gateway:`` server block (issue #3050).
|
||||
|
||||
@@ -896,7 +905,8 @@ class GatewayServerSchema(BaseModel):
|
||||
allowed_origins: Optional[List[str]] = None
|
||||
auth_token: Optional[str] = None
|
||||
auth: Optional[Dict[str, Any]] = None
|
||||
auth_scopes: Optional[Dict[str, List[str]]] = None
|
||||
# Issue #6253: token -> [scopes] or {scopes: [...], agents: [...]}.
|
||||
auth_scopes: Optional[Dict[str, Union[List[str], OperatorGrantSchema]]] = None
|
||||
max_connections: Optional[int] = Field(None, ge=0)
|
||||
max_sessions_per_agent: Optional[int] = Field(None, ge=0)
|
||||
session_config: Optional[Dict[str, Any]] = None
|
||||
|
||||
@@ -1427,6 +1427,36 @@
|
||||
"title": "HookSchema",
|
||||
"type": "object"
|
||||
},
|
||||
"OperatorGrantSchema": {
|
||||
"additionalProperties": false,
|
||||
"description": "Per-token operator grant (Issue #6253). Unknown keys are rejected so a\ntypo (``agnts``) cannot silently drop an agent confinement.",
|
||||
"properties": {
|
||||
"scopes": {
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"title": "Scopes",
|
||||
"type": "array"
|
||||
},
|
||||
"agents": {
|
||||
"anyOf": [
|
||||
{
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"title": "Agents"
|
||||
}
|
||||
},
|
||||
"title": "OperatorGrantSchema",
|
||||
"type": "object"
|
||||
},
|
||||
"OutboundHookSchema": {
|
||||
"additionalProperties": false,
|
||||
"description": "``gateway.hooks_outbound[]`` \u2014 a core ``CommandHook`` on lifecycle events (#6140).",
|
||||
@@ -1559,10 +1589,17 @@
|
||||
"anyOf": [
|
||||
{
|
||||
"additionalProperties": {
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
"anyOf": [
|
||||
{
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
{
|
||||
"$ref": "#/$defs/OperatorGrantSchema"
|
||||
}
|
||||
]
|
||||
},
|
||||
"type": "object"
|
||||
},
|
||||
|
||||
@@ -1287,6 +1287,9 @@ class WebSocketGateway:
|
||||
# {client_id: SessionSharingRole}). The owner stays ``session.client_id``.
|
||||
self._session_observers: Dict[str, Dict[str, Any]] = {}
|
||||
self._client_scopes: Dict[str, List[str]] = {} # client_id -> operator scopes
|
||||
# Issue #6253: client_id -> agent ids the operator is confined to
|
||||
# (absent/None = every agent).
|
||||
self._client_agents: Dict[str, Optional[List[str]]] = {}
|
||||
# Issue #3467: in-flight turn registry so a running turn can be aborted
|
||||
# (by a WS ``abort`` frame or a portable ``/stop`` chat command) and so
|
||||
# a per-turn timeout can cancel a runaway turn. Maps session_id ->
|
||||
@@ -2070,6 +2073,10 @@ class WebSocketGateway:
|
||||
self._register_client_conn(client_id, websocket)
|
||||
# Resolve operator scopes for this connection (all scopes if no policy).
|
||||
self._client_scopes[client_id] = self.config.resolve_scopes(operator_token)
|
||||
_resolve_agents = getattr(self.config, "resolve_agents", None) # older core: unconfined
|
||||
self._client_agents[client_id] = (
|
||||
_resolve_agents(operator_token) if _resolve_agents else None
|
||||
)
|
||||
# Issue #2661: stamp the session with the active secret's generation
|
||||
# so a later rotation can force-close it (instant revocation).
|
||||
self._client_auth_generation[client_id] = self._auth_generation()
|
||||
@@ -2120,6 +2127,7 @@ class WebSocketGateway:
|
||||
self._clients.pop(client_id, None)
|
||||
await self._teardown_client_conn(client_id)
|
||||
self._client_scopes.pop(client_id, None)
|
||||
self._client_agents.pop(client_id, None)
|
||||
self._client_auth_generation.pop(client_id, None)
|
||||
self._release_client_session(client_id)
|
||||
|
||||
@@ -3992,6 +4000,7 @@ class WebSocketGateway:
|
||||
self._client_conns.clear()
|
||||
self._client_sessions.clear()
|
||||
self._client_scopes.clear()
|
||||
self._client_agents.clear()
|
||||
self._client_auth_generation.clear()
|
||||
|
||||
for session_id in list(self._sessions.keys()):
|
||||
@@ -4104,8 +4113,14 @@ class WebSocketGateway:
|
||||
if msg_type == "hello":
|
||||
agent_id = data.get("agent_id")
|
||||
|
||||
# Check if agent exists
|
||||
if not agent_id or agent_id not in self._agents:
|
||||
# Check if agent exists. An agent outside the operator's grant
|
||||
# (Issue #6253) is reported identically so its existence is not
|
||||
# disclosed to a confined operator.
|
||||
if (
|
||||
not agent_id
|
||||
or agent_id not in self._agents
|
||||
or not self._client_may_reach_agent(client_id, agent_id)
|
||||
):
|
||||
error = HelloError(
|
||||
code=ConnectErrorCode.AGENT_NOT_FOUND,
|
||||
message=f"Agent not found: {agent_id}",
|
||||
@@ -4311,7 +4326,11 @@ class WebSocketGateway:
|
||||
# Keep backward compatibility with old join message
|
||||
elif msg_type == "join":
|
||||
agent_id = data.get("agent_id")
|
||||
if agent_id and agent_id in self._agents:
|
||||
if (
|
||||
agent_id
|
||||
and agent_id in self._agents
|
||||
and self._client_may_reach_agent(client_id, agent_id)
|
||||
):
|
||||
# Protocol version negotiation with validation
|
||||
try:
|
||||
client_min_version = int(data.get("min_version", MIN_PROTOCOL_VERSION))
|
||||
@@ -4373,6 +4392,18 @@ class WebSocketGateway:
|
||||
})
|
||||
return
|
||||
|
||||
# Issue #6253: deny resuming another agent's live session
|
||||
# *before* the resume call touches it.
|
||||
existing = self.get_session(session_id) if isinstance(session_id, str) else None
|
||||
if existing is not None and not self._client_may_reach_agent(
|
||||
client_id, existing.agent_id
|
||||
):
|
||||
await self._send_to_client(client_id, {
|
||||
"type": "error",
|
||||
"message": "Session does not belong to the requested agent",
|
||||
})
|
||||
return
|
||||
|
||||
# Resume or create session
|
||||
session, replay_events = self.resume_or_create_session(
|
||||
session_id=session_id,
|
||||
@@ -4380,7 +4411,17 @@ class WebSocketGateway:
|
||||
client_id=client_id,
|
||||
since_cursor=since_cursor,
|
||||
)
|
||||
|
||||
|
||||
# Issue #6253: a resumed session_id must not smuggle a confined
|
||||
# operator onto another agent's session (covers sessions
|
||||
# rehydrated from the persistent store).
|
||||
if not self._client_may_reach_agent(client_id, session.agent_id):
|
||||
await self._send_to_client(client_id, {
|
||||
"type": "error",
|
||||
"message": "Session does not belong to the requested agent",
|
||||
})
|
||||
return
|
||||
|
||||
# Set negotiated protocol version for the session
|
||||
session._protocol_version = negotiated_version
|
||||
|
||||
@@ -4646,7 +4687,7 @@ class WebSocketGateway:
|
||||
|
||||
elif msg_type == "server_reply":
|
||||
# Issue #5351: a client's correlated answer to a server→client
|
||||
# interactive request (approval / choice / input).
|
||||
# interactive request (approval / choice / input / secret).
|
||||
request_id = data.get("request_id")
|
||||
value = data.get("value")
|
||||
if not request_id or not isinstance(value, str):
|
||||
@@ -4680,7 +4721,7 @@ class WebSocketGateway:
|
||||
return True
|
||||
# Answering an approval on behalf of the turn is an approval decision
|
||||
# and requires the APPROVALS scope (parity with approvals.resolve);
|
||||
# choice/input answers require WRITE like sending a message.
|
||||
# choice/input/secret answers require WRITE like sending a message.
|
||||
required_scope = (
|
||||
OperatorScope.APPROVALS if req_kind == "approval" else OperatorScope.WRITE
|
||||
)
|
||||
@@ -5372,6 +5413,7 @@ class WebSocketGateway:
|
||||
# Mirror the normal disconnect cleanup so evicted clients do not leave
|
||||
# stale scope/session entries that accumulate until shutdown.
|
||||
self._client_scopes.pop(client_id, None)
|
||||
self._client_agents.pop(client_id, None)
|
||||
self._client_auth_generation.pop(client_id, None)
|
||||
self._release_client_session(client_id)
|
||||
if ws is not None:
|
||||
@@ -5442,6 +5484,7 @@ class WebSocketGateway:
|
||||
ws = self._clients.pop(client_id, None)
|
||||
await self._teardown_client_conn(client_id)
|
||||
self._client_scopes.pop(client_id, None)
|
||||
self._client_agents.pop(client_id, None)
|
||||
self._client_auth_generation.pop(client_id, None)
|
||||
self._release_client_session(client_id)
|
||||
if ws is not None:
|
||||
@@ -5494,6 +5537,7 @@ class WebSocketGateway:
|
||||
ws = self._clients.pop(client_id, None)
|
||||
await self._teardown_client_conn(client_id)
|
||||
self._client_scopes.pop(client_id, None)
|
||||
self._client_agents.pop(client_id, None)
|
||||
self._client_auth_generation.pop(client_id, None)
|
||||
self._release_client_session(client_id)
|
||||
self._reaped_connections += 1
|
||||
@@ -5851,7 +5895,9 @@ class WebSocketGateway:
|
||||
|
||||
* ``approval`` → ``"allow"`` or ``"deny"``.
|
||||
* ``choice`` → one of the offered ``options`` (any value if none given).
|
||||
* ``input`` → any string.
|
||||
* ``input`` / ``secret`` → any string. A ``secret`` value is handed
|
||||
only to the waiting caller — never logged, echoed, or journalled
|
||||
(Issue #6223).
|
||||
"""
|
||||
kind = getattr(req, "kind", "input")
|
||||
if kind == "approval":
|
||||
@@ -6837,6 +6883,9 @@ class WebSocketGateway:
|
||||
role = SessionSharingRole(role)
|
||||
if role is SessionSharingRole.OWNER or client_id == session.client_id:
|
||||
raise ValueError("The owner is bound at session creation, not attached")
|
||||
# Issue #6253: a confined operator must not join another agent's session.
|
||||
if not self._client_may_reach_agent(client_id, session.agent_id):
|
||||
raise PermissionError(f"Client {client_id} may not reach this session's agent")
|
||||
visibility = SessionVisibility(session._visibility)
|
||||
if visibility is SessionVisibility.PRIVATE:
|
||||
raise PermissionError(f"Session {session_id} is private")
|
||||
@@ -6999,6 +7048,15 @@ class WebSocketGateway:
|
||||
return func
|
||||
return decorator
|
||||
|
||||
def _client_may_reach_agent(self, client_id: str, agent_id: Any) -> bool:
|
||||
"""Whether a connected client's operator grant covers ``agent_id``.
|
||||
|
||||
Issue #6253: unconfined clients (no ``agents`` in their grant) reach
|
||||
every agent, preserving prior behaviour.
|
||||
"""
|
||||
allowed = self._client_agents.get(client_id)
|
||||
return allowed is None or agent_id in allowed
|
||||
|
||||
def _client_has_scope(self, client_id: str, scope: OperatorScope) -> bool:
|
||||
"""Whether a connected client holds ``scope`` (ADMIN implies all).
|
||||
|
||||
@@ -7043,6 +7101,11 @@ class WebSocketGateway:
|
||||
|
||||
def _event_visible_to(self, event: GatewayEvent, client_id: str) -> bool:
|
||||
"""Whether ``event`` should be delivered to ``client_id`` given its scopes."""
|
||||
# Issue #6253: an event sourced from an agent outside the client's
|
||||
# operator grant is withheld; non-agent (system) events are unaffected.
|
||||
source = getattr(event, "source", None)
|
||||
if source in self._agents and not self._client_may_reach_agent(client_id, source):
|
||||
return False
|
||||
return self._client_has_scope(client_id, self._event_required_scope(event))
|
||||
|
||||
async def emit(self, event: GatewayEvent) -> None:
|
||||
@@ -12060,6 +12123,42 @@ class WebSocketGateway:
|
||||
self.config.trusted_proxies = _trusted
|
||||
except Exception: # pragma: no cover - config is always a GatewayConfig
|
||||
pass
|
||||
# Issue #6253: apply the YAML operator policy (scopes and agent
|
||||
# confinement). Only overwrite when the YAML declares one, so a policy
|
||||
# a Python caller set on ``GatewayConfig`` survives.
|
||||
try:
|
||||
from praisonaiagents.gateway.config import _parse_auth_scopes
|
||||
except ImportError: # older core without per-agent grants
|
||||
_parse_auth_scopes = None
|
||||
_a = gw_cfg.get("auth")
|
||||
if gw_cfg.get("auth_scopes") or (isinstance(_a, dict) and _a.get("tokens")):
|
||||
# A declared policy that cannot be enforced must not start
|
||||
# with every operator unconfined (fail closed).
|
||||
raise ValueError(
|
||||
"gateway operator policy is declared but the installed "
|
||||
"praisonaiagents cannot enforce it; upgrade praisonaiagents"
|
||||
)
|
||||
_yaml_scopes = _parse_auth_scopes(gw_cfg) if _parse_auth_scopes else None
|
||||
if _yaml_scopes is not None:
|
||||
_auth = gw_cfg.get("auth")
|
||||
_tokens_form = isinstance(_auth, dict) and isinstance(_auth.get("tokens"), list)
|
||||
# ``load_gateway_config`` resolves ``${VAR}`` in values only, so
|
||||
# ``auth_scopes`` *keys* still hold placeholders here. An unset
|
||||
# var resolves to "" and that entry is dropped (it matches no one).
|
||||
_policy = {}
|
||||
for _tok, _grant in _yaml_scopes.items():
|
||||
_tok = _tok if _tokens_form else self._substitute_env_vars(_tok)
|
||||
if _tok:
|
||||
_policy[_tok] = _grant
|
||||
_declared = _auth.get("tokens") if _tokens_form else gw_cfg.get("auth_scopes")
|
||||
if _declared and not _policy:
|
||||
# An empty policy means "no policy" (every scope, every agent),
|
||||
# so refuse to start rather than fail open.
|
||||
raise ValueError(
|
||||
"gateway operator policy: every token resolved to empty "
|
||||
"(unset ${VAR}?); refusing to start without it"
|
||||
)
|
||||
self.config.auth_scopes = _policy
|
||||
# Issue #3593: the multi-bot CLI builds this gateway with a *default*
|
||||
# session config, so ``__init__`` already picked a store before this
|
||||
# YAML loaded. Re-read the ``session:`` block here and re-select the
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-bot"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
description = "Bots, gateway, and channel CLI for PraisonAI — messaging platforms and WebSocket control plane extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
||||
@@ -0,0 +1,284 @@
|
||||
"""Operator agent confinement on the gateway (Issue #6253).
|
||||
|
||||
``auth_scopes`` may carry ``{"scopes": [...], "agents": [...]}`` per token so
|
||||
an authenticated operator is confined to specific agents; the plain
|
||||
``token -> [scopes]`` form keeps today's gateway-wide behaviour.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.gateway import GatewayConfig
|
||||
from praisonaiagents.gateway.protocols import EventType, GatewayEvent
|
||||
from praisonai_bot.gateway.server import WebSocketGateway
|
||||
|
||||
|
||||
class _WS:
|
||||
def __init__(self):
|
||||
self.sent = []
|
||||
|
||||
async def send_json(self, data):
|
||||
self.sent.append(data)
|
||||
|
||||
async def close(self, code=None, reason=None):
|
||||
pass
|
||||
|
||||
|
||||
POLICY = {
|
||||
"tok_support": {"scopes": ["read", "write"], "agents": ["support"]},
|
||||
"tok_legacy": ["read", "write"],
|
||||
}
|
||||
|
||||
|
||||
def test_resolve_agents_contract():
|
||||
cfg = GatewayConfig(auth_token="t", auth_scopes=POLICY)
|
||||
assert cfg.resolve_scopes("tok_support") == ["read", "write"]
|
||||
assert cfg.resolve_agents("tok_support") == ["support"]
|
||||
# Plain list form and no-policy are unconfined (backward compatible).
|
||||
assert cfg.resolve_agents("tok_legacy") is None
|
||||
assert GatewayConfig().resolve_agents("anything") is None
|
||||
# Unknown token still holds no scopes.
|
||||
assert cfg.resolve_scopes("nope") == []
|
||||
|
||||
|
||||
def test_yaml_both_shapes_parse_agents():
|
||||
from praisonaiagents.gateway.config import MultiChannelGatewayConfig
|
||||
|
||||
flat = {"gateway": {"auth_scopes": {
|
||||
"tok_support": {"scopes": ["read", "write"], "agents": ["support"]},
|
||||
"tok_legacy": ["read"],
|
||||
}}}
|
||||
gw = MultiChannelGatewayConfig.from_dict(flat).gateway
|
||||
assert gw.resolve_agents("tok_support") == ["support"]
|
||||
assert gw.resolve_scopes("tok_support") == ["read", "write"]
|
||||
assert gw.resolve_agents("tok_legacy") is None
|
||||
|
||||
tokens = {"gateway": {"auth": {"tokens": [
|
||||
{"token": "tok_fin", "scopes": ["read"], "agents": ["finance"]},
|
||||
]}}}
|
||||
gw = MultiChannelGatewayConfig.from_dict(tokens).gateway
|
||||
assert gw.resolve_agents("tok_fin") == ["finance"]
|
||||
assert gw.resolve_scopes("tok_fin") == ["read"]
|
||||
|
||||
|
||||
def _gateway_with_client(agents):
|
||||
gw = WebSocketGateway(config=GatewayConfig(auth_token="tmp"))
|
||||
gw._agents["support"] = object()
|
||||
gw._agents["finance"] = object()
|
||||
ws = _WS()
|
||||
gw._clients["c1"] = ws
|
||||
gw._client_scopes["c1"] = ["read", "write"]
|
||||
gw._client_agents["c1"] = agents
|
||||
return gw, ws
|
||||
|
||||
|
||||
@pytest.mark.parametrize("frame", ["hello", "join"])
|
||||
def test_confined_operator_cannot_bind_other_agent(frame):
|
||||
gw, ws = _gateway_with_client(["support"])
|
||||
asyncio.run(gw._handle_client_message("c1", {"type": frame, "agent_id": "finance"}))
|
||||
assert "c1" not in gw._client_sessions
|
||||
assert "not found" in str(ws.sent[-1]).lower()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("frame", ["hello", "join"])
|
||||
def test_confined_operator_binds_allowed_agent(frame):
|
||||
gw, _ = _gateway_with_client(["support"])
|
||||
asyncio.run(gw._handle_client_message("c1", {"type": frame, "agent_id": "support"}))
|
||||
assert gw._sessions[gw._client_sessions["c1"]].agent_id == "support"
|
||||
|
||||
|
||||
def test_join_cannot_resume_foreign_agent_session():
|
||||
gw, ws = _gateway_with_client(["support"])
|
||||
foreign = gw.create_session("finance", client_id="other")
|
||||
asyncio.run(gw._handle_client_message(
|
||||
"c1",
|
||||
{"type": "join", "agent_id": "support", "session_id": foreign.session_id},
|
||||
))
|
||||
assert "c1" not in gw._client_sessions
|
||||
assert ws.sent[-1]["type"] == "error"
|
||||
|
||||
|
||||
def test_join_denies_foreign_session_before_resume(monkeypatch):
|
||||
# The denial must happen before resume touches the foreign session.
|
||||
gw, ws = _gateway_with_client(["support"])
|
||||
foreign = gw.create_session("finance", client_id="other")
|
||||
calls = []
|
||||
monkeypatch.setattr(gw, "resume_or_create_session", lambda **kw: calls.append(kw))
|
||||
asyncio.run(gw._handle_client_message(
|
||||
"c1",
|
||||
{"type": "join", "agent_id": "support", "session_id": foreign.session_id, "since": 0},
|
||||
))
|
||||
assert calls == []
|
||||
assert ws.sent[-1]["type"] == "error"
|
||||
assert foreign.client_id == "other"
|
||||
|
||||
|
||||
def test_attach_refuses_confined_client_on_foreign_agent_session():
|
||||
gw, _ = _gateway_with_client(["support"])
|
||||
foreign = gw.create_session("finance", client_id="other", visibility="shared")
|
||||
with pytest.raises(PermissionError):
|
||||
gw.attach(foreign.session_id, "c1", role="member")
|
||||
assert "c1" not in gw._client_sessions
|
||||
assert "c1" not in gw._session_observers.get(foreign.session_id, {})
|
||||
own = gw.create_session("support", client_id="other", visibility="shared")
|
||||
gw.attach(own.session_id, "c1", role="member")
|
||||
assert gw._client_sessions["c1"] == own.session_id
|
||||
|
||||
|
||||
def test_unconfined_operator_reaches_every_agent():
|
||||
gw, _ = _gateway_with_client(None)
|
||||
asyncio.run(gw._handle_client_message("c1", {"type": "hello", "agent_id": "finance"}))
|
||||
assert gw._sessions[gw._client_sessions["c1"]].agent_id == "finance"
|
||||
|
||||
|
||||
def test_broadcast_withholds_other_agents_events():
|
||||
gw, ws = _gateway_with_client(["support"])
|
||||
|
||||
async def _run():
|
||||
await gw.broadcast(GatewayEvent(type=EventType.MESSAGE, data={}, source="finance"))
|
||||
await gw.broadcast(GatewayEvent(type=EventType.MESSAGE, data={}, source="support"))
|
||||
await gw.broadcast(GatewayEvent(type=EventType.MESSAGE, data={}, source="system"))
|
||||
|
||||
asyncio.run(_run())
|
||||
assert [f.get("source") for f in ws.sent] == ["support", "system"]
|
||||
|
||||
|
||||
def test_misspelled_grant_key_fails_closed():
|
||||
# A typo must not silently drop the confinement (it would reach every agent).
|
||||
from praisonaiagents.gateway.config import (
|
||||
ConfigValidationError, MultiChannelGatewayConfig,
|
||||
)
|
||||
from praisonai_bot.bots._config_schema import GatewayServerSchema
|
||||
|
||||
bad = {"tok": {"scopes": ["read"], "agnts": ["support"]}}
|
||||
with pytest.raises(ConfigValidationError):
|
||||
MultiChannelGatewayConfig.from_dict({"gateway": {"auth_scopes": bad}})
|
||||
with pytest.raises(Exception):
|
||||
GatewayServerSchema(auth_scopes=bad)
|
||||
GatewayServerSchema(auth_scopes=POLICY) # valid shapes still accepted
|
||||
|
||||
# Same for the ``auth.tokens[]`` form.
|
||||
bad_tokens = {"tokens": [{"token": "tok", "scopes": ["read"], "agnts": ["support"]}]}
|
||||
with pytest.raises(ConfigValidationError):
|
||||
MultiChannelGatewayConfig.from_dict({"gateway": {"auth": bad_tokens}})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gw", [
|
||||
{"auth_scopes": {"tok": {"scopes": ["read"], "agents": "support"}}},
|
||||
{"auth": {"tokens": [{"token": "tok", "scopes": ["read"], "agents": "support"}]}},
|
||||
])
|
||||
def test_scalar_agents_value_rejected(gw):
|
||||
# A scalar must not be split into single-character agent ids.
|
||||
from praisonaiagents.gateway.config import MultiChannelGatewayConfig
|
||||
|
||||
with pytest.raises(ValueError, match="agents"):
|
||||
MultiChannelGatewayConfig.from_dict({"gateway": gw})
|
||||
|
||||
|
||||
class _StopStartup(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gw_yaml", [
|
||||
{"auth_scopes": {"tok_support": {"scopes": ["read"], "agents": ["support"]}}},
|
||||
{"auth": {"tokens": [{"token": "tok_support", "scopes": ["read"], "agents": ["support"]}]}},
|
||||
])
|
||||
def test_start_with_config_applies_yaml_policy(gw_yaml, monkeypatch):
|
||||
# The YAML startup path must stamp the policy onto self.config, otherwise
|
||||
# resolve_agents() returns None and confinement is silently off.
|
||||
gw = WebSocketGateway(config=GatewayConfig(auth_token="tmp"))
|
||||
cfg = {"gateway": dict(gw_yaml, session={"persist": False})}
|
||||
monkeypatch.setattr(gw, "load_gateway_config", lambda path: cfg)
|
||||
monkeypatch.setattr(gw, "_session_store_explicit", False, raising=False)
|
||||
|
||||
def _stop(*a, **k):
|
||||
raise _StopStartup
|
||||
|
||||
monkeypatch.setattr(gw, "_build_session_store", _stop)
|
||||
with pytest.raises(_StopStartup):
|
||||
asyncio.run(gw.start_with_config("gateway.yaml"))
|
||||
assert gw.config.resolve_agents("tok_support") == ["support"]
|
||||
assert gw.config.resolve_scopes("tok_support") == ["read"]
|
||||
|
||||
|
||||
def test_start_with_config_keeps_python_policy_without_yaml_policy(monkeypatch):
|
||||
gw = WebSocketGateway(config=GatewayConfig(auth_token="tmp", auth_scopes=POLICY))
|
||||
monkeypatch.setattr(gw, "load_gateway_config", lambda path: {"gateway": {"session": {}}})
|
||||
monkeypatch.setattr(gw, "_session_store_explicit", False, raising=False)
|
||||
monkeypatch.setattr(gw, "_build_session_store", lambda *a, **k: (_ for _ in ()).throw(_StopStartup()))
|
||||
with pytest.raises(_StopStartup):
|
||||
asyncio.run(gw.start_with_config("gateway.yaml"))
|
||||
assert gw.config.auth_scopes == POLICY
|
||||
|
||||
|
||||
def _start(gw, cfg, monkeypatch):
|
||||
cfg["gateway"].setdefault("session", {"persist": False})
|
||||
monkeypatch.setattr(gw, "load_gateway_config", lambda path: cfg)
|
||||
monkeypatch.setattr(gw, "_session_store_explicit", False, raising=False)
|
||||
monkeypatch.setattr(gw, "_build_session_store", lambda *a, **k: (_ for _ in ()).throw(_StopStartup()))
|
||||
asyncio.run(gw.start_with_config("gateway.yaml"))
|
||||
|
||||
|
||||
def test_start_with_config_expands_env_placeholder_keys(monkeypatch):
|
||||
# The loader resolves ${VAR} in values only; auth_scopes keys must be
|
||||
# expanded too, and the literal placeholder must not be a grant key.
|
||||
monkeypatch.setenv("SUPPORT_TOKEN", "real-support-token")
|
||||
gw = WebSocketGateway(config=GatewayConfig(auth_token="tmp"))
|
||||
cfg = {"gateway": {"auth_scopes": {
|
||||
"${SUPPORT_TOKEN}": {"scopes": ["read"], "agents": ["support"]},
|
||||
"${UNSET_TOKEN_6253}": ["read", "write"],
|
||||
}}}
|
||||
with pytest.raises(_StopStartup):
|
||||
_start(gw, cfg, monkeypatch)
|
||||
assert gw.config.auth_scopes == {
|
||||
"real-support-token": {"scopes": ["read"], "agents": ["support"]},
|
||||
}
|
||||
assert gw.config.resolve_agents("real-support-token") == ["support"]
|
||||
assert gw.config.resolve_scopes("${SUPPORT_TOKEN}") == []
|
||||
assert gw.config.resolve_scopes("") == []
|
||||
|
||||
|
||||
def test_start_with_config_refuses_policy_with_only_unset_tokens(monkeypatch):
|
||||
# An all-empty policy would mean "no policy" (fail open), so startup stops.
|
||||
monkeypatch.delenv("UNSET_TOKEN_6253", raising=False)
|
||||
gw = WebSocketGateway(config=GatewayConfig(auth_token="tmp"))
|
||||
cfg = {"gateway": {"auth_scopes": {
|
||||
"${UNSET_TOKEN_6253}": {"scopes": ["read"], "agents": ["support"]},
|
||||
}}}
|
||||
with pytest.raises(ValueError, match="resolved to empty"):
|
||||
_start(gw, cfg, monkeypatch)
|
||||
# Tokens form: the loader already turned the unset value into "".
|
||||
cfg = {"gateway": {"auth": {"tokens": [
|
||||
{"token": "", "scopes": ["read"], "agents": ["support"]},
|
||||
]}}}
|
||||
with pytest.raises(ValueError, match="resolved to empty"):
|
||||
_start(gw, cfg, monkeypatch)
|
||||
|
||||
|
||||
def test_start_with_config_tolerates_core_without_policy_parser(monkeypatch):
|
||||
# An older praisonaiagents lacks _parse_auth_scopes / resolve_agents; with
|
||||
# no YAML policy the gateway must still start instead of ImportError.
|
||||
import praisonaiagents.gateway.config as core_cfg
|
||||
|
||||
monkeypatch.delattr(core_cfg, "_parse_auth_scopes")
|
||||
gw = WebSocketGateway(config=GatewayConfig(auth_token="tmp", auth_scopes=POLICY))
|
||||
with pytest.raises(_StopStartup):
|
||||
_start(gw, {"gateway": {}}, monkeypatch)
|
||||
assert gw.config.auth_scopes == POLICY
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gw_yaml", [
|
||||
{"auth_scopes": {"support-token": {"scopes": ["read"], "agents": ["support"]}}},
|
||||
{"auth": {"tokens": [{"token": "t", "scopes": ["read"], "agents": ["support"]}]}},
|
||||
])
|
||||
def test_start_with_config_refuses_unenforceable_policy_on_older_core(gw_yaml, monkeypatch):
|
||||
# A declared policy that the installed core cannot parse must fail closed,
|
||||
# not start with every operator unconfined.
|
||||
import praisonaiagents.gateway.config as core_cfg
|
||||
|
||||
monkeypatch.delattr(core_cfg, "_parse_auth_scopes")
|
||||
gw = WebSocketGateway(config=GatewayConfig(auth_token="tmp"))
|
||||
with pytest.raises(ValueError, match="cannot enforce"):
|
||||
_start(gw, {"gateway": dict(gw_yaml)}, monkeypatch)
|
||||
assert not gw.config.auth_scopes
|
||||
@@ -67,6 +67,37 @@ def test_server_reply_as_dict():
|
||||
}
|
||||
|
||||
|
||||
def test_server_reply_repr_hides_value():
|
||||
reply = GatewayServerReply(request_id="r1", value="hunter2")
|
||||
assert "hunter2" not in repr(reply)
|
||||
assert reply.value == "hunter2"
|
||||
|
||||
|
||||
def test_secret_request_resolves_only_to_caller():
|
||||
"""Issue #6223: a secret reply reaches the waiting caller and nowhere else."""
|
||||
async def go():
|
||||
gw = WebSocketGateway()
|
||||
sent = _connect(gw, "c1", "sess-1")
|
||||
req = GatewayServerRequest(
|
||||
request_id="req-s", kind="secret", prompt="sudo password?",
|
||||
session_id="sess-1",
|
||||
)
|
||||
task = asyncio.create_task(gw.request_channel.request(req, timeout_s=5))
|
||||
await asyncio.sleep(0)
|
||||
assert sent[0]["request"]["kind"] == "secret"
|
||||
|
||||
await gw._handle_client_message(
|
||||
"c1", {"type": "server_reply", "request_id": "req-s", "value": "hunter2"}
|
||||
)
|
||||
reply = await asyncio.wait_for(task, timeout=1)
|
||||
assert reply.value == "hunter2"
|
||||
# The value is never echoed back over the wire (ack carries no value).
|
||||
assert all("hunter2" not in repr(frame) for frame in sent)
|
||||
assert gw.request_channel.open_requests("sess-1") == []
|
||||
|
||||
asyncio.run(go())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test harness: a gateway with a connected, session-bound client that records
|
||||
# every frame the server sends it.
|
||||
|
||||
Generated
+2
-2
@@ -1428,7 +1428,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-bot"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -1510,7 +1510,7 @@ provides-extras = ["gateway", "bot", "bot-whatsapp-web", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.21"
|
||||
version = "1.7.22"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.20"
|
||||
__version__ = "0.1.21"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-browser"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
description = "Browser automation for PraisonAI — Chrome extension bridge, CDP, and Playwright agents extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
Generated
+2
-2
@@ -932,7 +932,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-browser"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -976,7 +976,7 @@ provides-extras = ["server", "playwright", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.21"
|
||||
version = "1.7.22"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -12,6 +12,6 @@ the remaining terminal-agent modules, with PEP 562 shims left behind at the
|
||||
old ``praisonai.*`` paths for backward compatibility.
|
||||
"""
|
||||
|
||||
__version__ = "0.1.20"
|
||||
__version__ = "0.1.21"
|
||||
|
||||
__all__ = ["__version__"]
|
||||
@@ -1712,6 +1712,7 @@ def run_main(
|
||||
# execution path does not discover and execute the same definitions twice.
|
||||
preloaded_agent_config = None
|
||||
custom_agent_model = None
|
||||
custom_agent_base_url = None
|
||||
if target or agent or command: # Any execution surface may need model/budget resolution
|
||||
import sys
|
||||
from praisonai_code.llm.credentials import ensure_configured_or_onboard
|
||||
@@ -1731,6 +1732,10 @@ def run_main(
|
||||
_preview_llm = (_preview_config or {}).get("llm")
|
||||
if isinstance(_preview_llm, dict):
|
||||
custom_agent_model = _preview_llm.get("model")
|
||||
# Gate an Ollama agent on the server it will really call.
|
||||
_base = _preview_llm.get("base_url")
|
||||
if isinstance(_base, str) and _base.strip():
|
||||
custom_agent_base_url = _base.strip()
|
||||
else:
|
||||
custom_agent_model = _preview_llm
|
||||
if isinstance(custom_agent_model, str):
|
||||
@@ -1747,7 +1752,9 @@ def run_main(
|
||||
# endpoint, credentials, sampling, and retry options.
|
||||
credential_model = model or custom_agent_model
|
||||
resolved_model = ensure_configured_or_onboard(
|
||||
model=credential_model, interactive=not _headless
|
||||
model=credential_model,
|
||||
interactive=not _headless,
|
||||
base_url=None if model else custom_agent_base_url,
|
||||
)
|
||||
if not (agent and model is None and custom_agent_model):
|
||||
model = resolved_model
|
||||
@@ -2386,6 +2393,8 @@ def _run_prompt(
|
||||
if allow_local_tools:
|
||||
_os.environ[_ALLOW_LOCAL_TOOLS_ENV] = "true"
|
||||
|
||||
_manifest_session = None
|
||||
|
||||
try:
|
||||
# Handle session continuity first (before any execution mode)
|
||||
from praisonai_code.cli.main import PraisonAI
|
||||
@@ -2435,6 +2444,13 @@ def _run_prompt(
|
||||
import uuid
|
||||
auto_save_name = session_id or "session-" + str(uuid.uuid4())[:8]
|
||||
|
||||
# Write-ahead "run in flight" marker on the session manifest, released
|
||||
# in ``finally`` below; a run that dies before then is reported as
|
||||
# interrupted on the next --continue/--session (Issue #6248).
|
||||
_manifest_session = _begin_manifest_run(
|
||||
session_id or auto_save_name, prompt, output, quiet=_json_stdout
|
||||
)
|
||||
|
||||
# Detect-and-attach: if a warm runtime is running, forward the plain
|
||||
# prompt to it (skipping per-invocation cold-start) and fall back to
|
||||
# in-process execution otherwise. Only the simple text path attaches;
|
||||
@@ -2767,6 +2783,7 @@ def _run_prompt(
|
||||
output.print_error(str(e))
|
||||
raise typer.Exit(1)
|
||||
finally:
|
||||
_end_manifest_run(_manifest_session)
|
||||
# Restore the prior opt-in state so the grant is strictly per-invocation.
|
||||
if _prev_allow_local_tools is None:
|
||||
_os.environ.pop(_ALLOW_LOCAL_TOOLS_ENV, None)
|
||||
@@ -2778,6 +2795,57 @@ def _run_prompt(
|
||||
output.mode = _prev_mode
|
||||
|
||||
|
||||
def _begin_manifest_run(session_key, prompt, output, *, quiet: bool = False):
|
||||
"""Commit the session manifest's in-flight marker for this run.
|
||||
|
||||
If the previous run on this session never released its marker it was
|
||||
interrupted; say so, so the user can re-issue it. Returns the session key
|
||||
a ``(manifest, session_key, run_id)`` handle to release later, or ``None``
|
||||
when there is no session or the manifest is unavailable (manifest problems
|
||||
never block a run).
|
||||
"""
|
||||
if not session_key:
|
||||
return None
|
||||
try:
|
||||
import uuid
|
||||
from praisonaiagents.session.manifest import FileSessionManifest
|
||||
from ..utils.project import get_project_sessions_dir
|
||||
|
||||
# Scope manifests like the project session store, so equal session
|
||||
# names in different projects never share a run marker.
|
||||
manifest = FileSessionManifest(str(get_project_sessions_dir() / "manifests"))
|
||||
stale = manifest.interrupted(session_key)
|
||||
if stale is not None and not quiet:
|
||||
# The prompt itself is not stored (only a digest), so it is not echoed.
|
||||
output.print_warning(
|
||||
f"The previous run in this session was interrupted "
|
||||
f"(attempt {stale.attempt}); re-issue it to retry."
|
||||
)
|
||||
run_id = uuid.uuid4().hex[:12]
|
||||
manifest.begin_run(session_key, run_id, run_input=prompt)
|
||||
return manifest, session_key, run_id
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _end_manifest_run(handle) -> None:
|
||||
"""Release this run's own in-flight marker unless it was interrupted (Ctrl+C)."""
|
||||
if not handle:
|
||||
return
|
||||
import sys
|
||||
|
||||
exc = sys.exc_info()[1]
|
||||
if isinstance(exc, KeyboardInterrupt):
|
||||
return
|
||||
manifest, session_key, run_id = handle
|
||||
try:
|
||||
manifest.end_run(
|
||||
session_key, "completed" if exc is None else "ended", run_id=run_id
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _record_session_usage(session_id, model, output, *, suppress_info: bool = False) -> None:
|
||||
"""Accumulate this run's token/cost usage into the active session and show
|
||||
a compact running total footer (Issue #2421).
|
||||
|
||||
@@ -123,6 +123,43 @@ def detect_local_endpoint():
|
||||
return None
|
||||
|
||||
|
||||
_OLLAMA_PREFIXES = ("ollama/", "ollama_chat/")
|
||||
|
||||
|
||||
def _ollama_host_for_model(
|
||||
model: Optional[str], base_url: Optional[str] = None
|
||||
) -> Optional[str]:
|
||||
"""Return the Ollama host to probe for an ``ollama/`` model, else ``None``.
|
||||
|
||||
Ollama needs no key, so an explicit Ollama model is gated on the server
|
||||
being reachable rather than on a credential. The endpoint the caller will
|
||||
actually use (e.g. an agent's ``llm.base_url``) wins, then
|
||||
``OLLAMA_API_BASE`` (what the runtime honours), then ``OLLAMA_HOST``, then
|
||||
the default port.
|
||||
"""
|
||||
import os
|
||||
|
||||
if not model or not model.lower().startswith(_OLLAMA_PREFIXES):
|
||||
return None
|
||||
from praisonai_code.llm.local_detect import _DEFAULT_OLLAMA_HOST
|
||||
|
||||
return (
|
||||
base_url
|
||||
or os.environ.get("OLLAMA_API_BASE")
|
||||
or os.environ.get("OLLAMA_HOST")
|
||||
or _DEFAULT_OLLAMA_HOST
|
||||
)
|
||||
|
||||
|
||||
def _ollama_reachable(host: str) -> bool:
|
||||
try:
|
||||
from praisonai_code.llm.local_detect import ollama_reachable
|
||||
|
||||
return ollama_reachable(host)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _provider_key_vars_for_model(model: str) -> tuple[str, ...]:
|
||||
"""Map a model id to the environment variable(s) for its provider.
|
||||
|
||||
@@ -258,6 +295,10 @@ def is_configured(model: Optional[str] = None) -> bool:
|
||||
if model is None:
|
||||
model = default_model_for_available_provider()
|
||||
|
||||
ollama_host = _ollama_host_for_model(model) if explicit_model else None
|
||||
if ollama_host is not None:
|
||||
return _ollama_reachable(ollama_host)
|
||||
|
||||
required_vars = _provider_key_vars_for_model(model)
|
||||
|
||||
if required_vars:
|
||||
@@ -303,6 +344,7 @@ def ensure_configured_or_onboard(
|
||||
*,
|
||||
model: Optional[str] = None,
|
||||
interactive: bool = True,
|
||||
base_url: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
"""Shared first-run credential gate for LLM-bound interactive commands.
|
||||
|
||||
@@ -331,6 +373,9 @@ def ensure_configured_or_onboard(
|
||||
keyless local-first fallback (its own provider gate applies).
|
||||
interactive: Whether a wizard prompt may be shown. Callers pass ``False``
|
||||
for headless (``-p``, piped, ``--output json``) paths.
|
||||
base_url: The endpoint the caller will actually use (e.g. a custom
|
||||
agent's ``llm.base_url``). For an explicit Ollama model this is the
|
||||
server probed, instead of ``OLLAMA_API_BASE``/``OLLAMA_HOST``.
|
||||
|
||||
Returns:
|
||||
The model id to proceed with. May be a detected local model id when the
|
||||
@@ -343,6 +388,21 @@ def ensure_configured_or_onboard(
|
||||
import typer
|
||||
|
||||
inject_credentials_into_env()
|
||||
|
||||
# An explicit Ollama model needs a running server, not a key: probe the
|
||||
# endpoint the caller will really use and, if it is down, say so instead
|
||||
# of sending the user to an API-key wizard.
|
||||
ollama_host = _ollama_host_for_model(model, base_url)
|
||||
if ollama_host is not None:
|
||||
if _ollama_reachable(ollama_host):
|
||||
return model
|
||||
typer.echo(
|
||||
f"Error: Ollama isn't reachable at {ollama_host} — start it with "
|
||||
"`ollama serve` (or set OLLAMA_API_BASE).",
|
||||
err=True,
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
if is_configured(model):
|
||||
return model
|
||||
|
||||
@@ -360,11 +420,16 @@ def ensure_configured_or_onboard(
|
||||
return local.model
|
||||
|
||||
if not interactive:
|
||||
# Only mention auto-detection when it actually ran and found nothing.
|
||||
hint = f" for model {model}" if model else ""
|
||||
tail = (
|
||||
"" if model else
|
||||
"\n(no running local endpoint such as Ollama was found; one "
|
||||
"would be used automatically)"
|
||||
)
|
||||
typer.echo(
|
||||
"Error: No API key configured. Run: praisonai setup\n"
|
||||
"or set environment variables like OPENAI_API_KEY\n"
|
||||
"(a running local endpoint such as Ollama would be used "
|
||||
"automatically)",
|
||||
f"Error: No API key configured{hint}. Run: praisonai setup\n"
|
||||
f"or set environment variables like OPENAI_API_KEY{tail}",
|
||||
err=True,
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
@@ -19,9 +19,10 @@ from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
# Total time budget for a probe so the credential/first-run hot path is never
|
||||
# stalled when nothing is listening. Kept small deliberately.
|
||||
_PROBE_TIMEOUT_S = 0.15
|
||||
# Per-request probe timeout. Only paid on the keyless path; a closed port is
|
||||
# refused instantly, while a running Ollama can take ~0.35 s on its first call,
|
||||
# so a tighter budget made detection flaky (issue #6228).
|
||||
_PROBE_TIMEOUT_S = 1.0
|
||||
|
||||
# How long a negative probe is remembered so repeated invocations in the same
|
||||
# process (e.g. bare TUI then run) don't re-pay the connection latency.
|
||||
@@ -187,6 +188,11 @@ def _probe_ollama_tags(host: str) -> Optional[str]:
|
||||
return None
|
||||
|
||||
|
||||
def ollama_reachable(host: str) -> bool:
|
||||
"""Return ``True`` if an Ollama server answers ``/api/tags`` at ``host``."""
|
||||
return _get_json(_root_host(_with_scheme(host)) + "/api/tags") is not None
|
||||
|
||||
|
||||
def detect_local_model(*, use_cache: bool = True) -> Optional[LocalModel]:
|
||||
"""Detect a reachable local OpenAI-compatible endpoint.
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-code"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
description = "Agentic terminal CLI for PraisonAI — the terminal-native agent (run, chat, code, runtime, cli backends) extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
@@ -0,0 +1,144 @@
|
||||
"""Tests for issue #6228 — explicit ``ollama/`` models need no API key.
|
||||
|
||||
An explicit Ollama model is gated on the server being reachable, not on
|
||||
``OLLAMA_HOST`` being set; the probe tolerates a slow first response; and the
|
||||
error text for an explicit model does not claim auto-detection would kick in.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
|
||||
import pytest
|
||||
import typer
|
||||
|
||||
from praisonai_code.llm import credentials, local_detect
|
||||
|
||||
|
||||
_ENV = (
|
||||
"OPENAI_API_KEY", "ANTHROPIC_API_KEY", "GEMINI_API_KEY", "GOOGLE_API_KEY",
|
||||
"GROQ_API_KEY", "COHERE_API_KEY", "OPENROUTER_API_KEY", "OLLAMA_HOST",
|
||||
"OLLAMA_API_BASE", "OPENAI_BASE_URL", "MODEL_NAME", "OPENAI_MODEL_NAME",
|
||||
"PRAISONAI_LOCAL_ENDPOINTS",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clean_env(monkeypatch):
|
||||
for var in _ENV:
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
monkeypatch.setattr(credentials, "inject_credentials_into_env", lambda: False)
|
||||
local_detect.reset_cache()
|
||||
yield monkeypatch
|
||||
local_detect.reset_cache()
|
||||
|
||||
|
||||
class _Resp(io.BytesIO):
|
||||
status = 200
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
|
||||
def _fake_urlopen(seen, delay=0.0):
|
||||
def _urlopen(url, timeout):
|
||||
seen.append((url, timeout))
|
||||
if delay > timeout:
|
||||
raise TimeoutError("probe timed out")
|
||||
return _Resp(json.dumps({"models": [{"name": "llama3.2:latest"}]}).encode())
|
||||
return _urlopen
|
||||
|
||||
|
||||
def _refuse(url, timeout):
|
||||
raise ConnectionRefusedError()
|
||||
|
||||
|
||||
def test_explicit_ollama_configured_when_reachable(clean_env):
|
||||
seen = []
|
||||
clean_env.setattr(local_detect.urllib.request, "urlopen", _fake_urlopen(seen))
|
||||
assert credentials.is_configured("ollama/llama3.2") is True
|
||||
assert seen[0][0] == "http://127.0.0.1:11434/api/tags"
|
||||
|
||||
|
||||
def test_explicit_ollama_honours_ollama_api_base(clean_env):
|
||||
seen = []
|
||||
clean_env.setenv("OLLAMA_API_BASE", "http://gpu-box:11434")
|
||||
clean_env.setattr(local_detect.urllib.request, "urlopen", _fake_urlopen(seen))
|
||||
assert credentials.is_configured("ollama_chat/llama3.2") is True
|
||||
assert seen[0][0] == "http://gpu-box:11434/api/tags"
|
||||
|
||||
|
||||
def test_explicit_ollama_not_configured_when_unreachable(clean_env):
|
||||
clean_env.setattr(local_detect.urllib.request, "urlopen", _refuse)
|
||||
assert credentials.is_configured("ollama/llama3.2") is False
|
||||
|
||||
|
||||
def test_unreachable_ollama_error_is_actionable(clean_env, capsys):
|
||||
clean_env.setattr(local_detect.urllib.request, "urlopen", _refuse)
|
||||
with pytest.raises(typer.Exit):
|
||||
credentials.ensure_configured_or_onboard(
|
||||
model="ollama/llama3.2", interactive=False
|
||||
)
|
||||
err = capsys.readouterr().err
|
||||
assert "Ollama isn't reachable at http://127.0.0.1:11434" in err
|
||||
assert "ollama serve" in err
|
||||
assert "API key" not in err
|
||||
|
||||
|
||||
def test_explicit_ollama_passes_gate_when_reachable(clean_env):
|
||||
clean_env.setattr(local_detect.urllib.request, "urlopen", _fake_urlopen([]))
|
||||
assert credentials.ensure_configured_or_onboard(
|
||||
model="ollama/llama3.2", interactive=False
|
||||
) == "ollama/llama3.2"
|
||||
|
||||
|
||||
def test_probe_tolerates_slow_first_response(clean_env):
|
||||
"""A ~0.4 s first /api/tags response must still be detected."""
|
||||
clean_env.setattr(
|
||||
local_detect.urllib.request, "urlopen", _fake_urlopen([], delay=0.4)
|
||||
)
|
||||
result = local_detect.detect_local_model(use_cache=False)
|
||||
assert result is not None
|
||||
assert result.model == "ollama/llama3.2:latest"
|
||||
|
||||
|
||||
def test_explicit_cloud_model_error_omits_auto_detect_hint(clean_env, capsys):
|
||||
clean_env.setattr(credentials, "detect_local_endpoint", lambda: None)
|
||||
clean_env.setattr(credentials, "is_configured", lambda model=None: False)
|
||||
with pytest.raises(typer.Exit):
|
||||
credentials.ensure_configured_or_onboard(
|
||||
model="gpt-4o-mini", interactive=False
|
||||
)
|
||||
err = capsys.readouterr().err
|
||||
assert "for model gpt-4o-mini" in err
|
||||
assert "automatically" not in err
|
||||
|
||||
|
||||
def test_keyless_no_model_error_mentions_auto_detect(clean_env, capsys):
|
||||
clean_env.setattr(credentials, "detect_local_endpoint", lambda: None)
|
||||
clean_env.setattr(credentials, "is_configured", lambda model=None: False)
|
||||
with pytest.raises(typer.Exit):
|
||||
credentials.ensure_configured_or_onboard(model=None, interactive=False)
|
||||
assert "automatically" in capsys.readouterr().err
|
||||
|
||||
|
||||
def test_gate_probes_agent_base_url_not_env_host(clean_env):
|
||||
"""A custom agent's ``llm.base_url`` is the server gated, not OLLAMA_HOST."""
|
||||
seen = []
|
||||
clean_env.setenv("OLLAMA_HOST", "http://down-box:11434")
|
||||
|
||||
def _urlopen(url, timeout):
|
||||
seen.append(url)
|
||||
if "down-box" in url:
|
||||
raise ConnectionRefusedError()
|
||||
return _Resp(json.dumps({"models": []}).encode())
|
||||
|
||||
clean_env.setattr(local_detect.urllib.request, "urlopen", _urlopen)
|
||||
assert credentials.ensure_configured_or_onboard(
|
||||
model="ollama/llama3.2",
|
||||
interactive=False,
|
||||
base_url="http://gpu-box:11434/v1",
|
||||
) == "ollama/llama3.2"
|
||||
assert seen == ["http://gpu-box:11434/api/tags"]
|
||||
Generated
+2
-2
@@ -1370,7 +1370,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-code"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1403,7 +1403,7 @@ requires-dist = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.21"
|
||||
version = "1.7.22"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.20"
|
||||
__version__ = "0.1.21"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-deploy"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
description = "Deployment for PraisonAI — API, Docker, and cloud (AWS, Azure, GCP) deployments extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
Generated
+2
-2
@@ -977,7 +977,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-deploy"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1014,7 +1014,7 @@ provides-extras = ["api", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.21"
|
||||
version = "1.7.22"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.22"
|
||||
__version__ = "0.1.23"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-mcp"
|
||||
version = "0.1.22"
|
||||
version = "0.1.23"
|
||||
description = "MCP server hosting for PraisonAI — expose agents and tools to Cursor, Claude Desktop, and other MCP clients."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
Generated
+2
-2
@@ -1044,7 +1044,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-mcp"
|
||||
version = "0.1.22"
|
||||
version = "0.1.23"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1084,7 +1084,7 @@ provides-extras = ["server", "auth", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.21"
|
||||
version = "1.7.22"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.20"
|
||||
__version__ = "0.1.21"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-sandbox"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
description = "Where PraisonAI agents run code — Docker, E2B, Modal, Daytona, Fly.io, Tenki, Novita, Sandlock, SSH and subprocess, as both single sandboxes and multi-instance compute providers."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
Generated
+2
-2
@@ -1679,7 +1679,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-sandbox"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1744,7 +1744,7 @@ provides-extras = ["docker", "e2b", "sandlock", "ssh", "tenki", "flyio", "modal"
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.21"
|
||||
version = "1.7.22"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.20"
|
||||
__version__ = "0.1.21"
|
||||
@@ -717,6 +717,9 @@ def train_export(
|
||||
mtp_draft: bool = typer.Option(
|
||||
False, "--mtp-draft/--no-mtp-draft",
|
||||
help="Also download the stock MTP drafter (Gemma-4 only) for fast inference"),
|
||||
push: Optional[bool] = typer.Option(
|
||||
None, "--push/--no-push",
|
||||
help="ollama: push to ollama.com after `ollama create` (--no-push keeps it local)"),
|
||||
):
|
||||
"""
|
||||
Export / publish an ALREADY-trained model without re-running training.
|
||||
@@ -726,6 +729,7 @@ def train_export(
|
||||
praisonai-train export gguf --model-dir lora_model --quant q4_k_m # local .gguf
|
||||
praisonai-train export gguf --model-dir lora_model --hf me/my-model # local + push
|
||||
praisonai-train export ollama --model-dir lora_model --ollama me/my-model --quant q4_k_m
|
||||
praisonai-train export ollama --model-dir lora_model --ollama my-model --no-push # local only
|
||||
"""
|
||||
from ..output.console import get_output_controller
|
||||
|
||||
@@ -765,6 +769,9 @@ def train_export(
|
||||
cfg["hf_model_name"] = hf
|
||||
if ollama:
|
||||
cfg["ollama_model"] = ollama
|
||||
# Either explicit flag overrides the config; unset keeps config/default (push).
|
||||
if isinstance(push, bool):
|
||||
cfg["ollama_push"] = push
|
||||
|
||||
# Recover the model name (drives chat-template selection) from --base-model or
|
||||
# the trained model's config.json (_name_or_path), falling back to the dir name.
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
ollama_save: "true"
|
||||
ollama_push: "true" # set "false" to keep the Ollama model local (no upload to ollama.com)
|
||||
huggingface_save: "true"
|
||||
train: "true"
|
||||
|
||||
|
||||
@@ -220,8 +220,9 @@ def create_and_push_ollama_model(
|
||||
model_parameters: str,
|
||||
modelfile_content: str,
|
||||
quantization: Optional[str] = None,
|
||||
push: bool = True,
|
||||
) -> None:
|
||||
"""Create and push an Ollama model with proper daemon management.
|
||||
"""Create (and by default push) an Ollama model with proper daemon management.
|
||||
|
||||
Args:
|
||||
ollama_model: Name of the Ollama model
|
||||
@@ -229,6 +230,8 @@ def create_and_push_ollama_model(
|
||||
modelfile_content: Content for the Modelfile
|
||||
quantization: Optional quantization (e.g. "q4_k_m") passed to
|
||||
``ollama create --quantize`` so the model isn't stored as huge f16.
|
||||
push: When False, stop after a successful ``ollama create`` and keep the
|
||||
model local-only (nothing is uploaded to ollama.com).
|
||||
|
||||
Raises:
|
||||
RuntimeError: If ollama create/push operations fail (with actionable guidance)
|
||||
@@ -259,6 +262,10 @@ def create_and_push_ollama_model(
|
||||
f"{(create.stderr or create.stdout or '').strip()}"
|
||||
)
|
||||
|
||||
if not push:
|
||||
print(f"Created local Ollama model `{tag}` (not pushed). Run it with: ollama run {tag}")
|
||||
return
|
||||
|
||||
# Push with captured output so a 401/403 becomes an actionable message rather
|
||||
# than a raw non-zero exit.
|
||||
result = subprocess.run(["ollama", "push", tag], capture_output=True, text=True)
|
||||
|
||||
@@ -545,6 +545,7 @@ class TrainModel:
|
||||
"optim", "weight_decay", "lr_scheduler_type", "seed", "output_dir",
|
||||
"assistant_only_loss", "train_on_responses_only", "save_steps",
|
||||
"train", "huggingface_save", "huggingface_save_gguf", "ollama_save",
|
||||
"ollama_push",
|
||||
"method", "beta", "max_prompt_length", "desirable_weight", "undesirable_weight",
|
||||
"reward_funcs", "num_generations", "max_completion_length",
|
||||
"embedding_learning_rate",
|
||||
@@ -951,6 +952,10 @@ class TrainModel:
|
||||
if self._flag(dataset_info.get("shuffle"), default=False):
|
||||
dataset = dataset.shuffle(
|
||||
seed=int(dataset_info.get("seed", self.config.get("seed", 3407))))
|
||||
# OpenAI-style chat rows ({"messages": [...]}, as written by
|
||||
# praisonaiagents.eval.export_sft) are the ShareGPT shape under another name.
|
||||
if "messages" in dataset.column_names and "conversations" not in dataset.column_names:
|
||||
dataset = dataset.rename_column("messages", "conversations")
|
||||
if "conversations" in dataset.column_names:
|
||||
print("DEBUG: Standardizing dataset (ShareGPT style)...")
|
||||
dataset = standardize_sharegpt(dataset)
|
||||
@@ -970,8 +975,8 @@ class TrainModel:
|
||||
if len(dataset) == 0:
|
||||
raise ValueError(
|
||||
"All examples formatted to empty text. The dataset columns are "
|
||||
f"{columns_before}; this trainer understands ShareGPT "
|
||||
"('conversations'), Alpaca ('instruction'/'input'/'output') and "
|
||||
f"{columns_before}; this trainer understands chat "
|
||||
"('conversations' or 'messages'), Alpaca ('instruction'/'input'/'output') and "
|
||||
"raw text ('text'). Rename your columns to one of those shapes, "
|
||||
"or set chat_template to match your data.")
|
||||
return dataset
|
||||
@@ -1784,6 +1789,7 @@ class TrainModel:
|
||||
self.config.get('model_parameters', 'latest'),
|
||||
modelfile_content,
|
||||
quantization=self.config.get('quantization_method'),
|
||||
push=self._flag(self.config.get('ollama_push'), default=True),
|
||||
)
|
||||
|
||||
def run(self):
|
||||
|
||||
@@ -313,7 +313,8 @@ PARAMETER top_p 0.9
|
||||
create_and_push_ollama_model(
|
||||
self.config['ollama_model'],
|
||||
self.config['model_parameters'],
|
||||
modelfile_content
|
||||
modelfile_content,
|
||||
push=self._flag(self.config.get('ollama_push'), default=True),
|
||||
)
|
||||
|
||||
def run(self):
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-train"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
description = "Training for PraisonAI — LLM fine-tuning (Unsloth) and iterative agent training extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
"""An OpenAI-style `messages` chat dataset formatted to empty text.
|
||||
|
||||
`praisonai-train validate` and `praisonaiagents.eval.export_sft` both use the
|
||||
`{"messages": [...]}` shape, but `process_dataset` only recognised
|
||||
`conversations`, so every row fell into the Alpaca branch, formatted to "",
|
||||
and the run died after the base model had already been loaded.
|
||||
"""
|
||||
|
||||
from praisonai_train.train.llm import trainer as trainer_mod
|
||||
|
||||
|
||||
class _FakeDataset:
|
||||
def __init__(self, rows):
|
||||
self.rows = rows
|
||||
|
||||
@property
|
||||
def column_names(self):
|
||||
return list(self.rows[0].keys()) if self.rows else []
|
||||
|
||||
def __len__(self):
|
||||
return len(self.rows)
|
||||
|
||||
def rename_column(self, src, dst):
|
||||
return _FakeDataset([{(dst if k == src else k): v for k, v in r.items()}
|
||||
for r in self.rows])
|
||||
|
||||
def map(self, func, batched, remove_columns):
|
||||
batch = {c: [r[c] for r in self.rows] for c in self.column_names}
|
||||
out = func(batch)
|
||||
return _FakeDataset([{"text": t} for t in out["text"]])
|
||||
|
||||
def filter(self, pred):
|
||||
return _FakeDataset([r for r in self.rows if pred(r)])
|
||||
|
||||
|
||||
class _Tokenizer:
|
||||
def apply_chat_template(self, convo, tokenize, add_generation_prompt):
|
||||
return "".join(f"<{m['role']}>{m['content']}" for m in convo)
|
||||
|
||||
|
||||
def test_messages_column_is_treated_as_conversations(monkeypatch):
|
||||
rows = [{"messages": [
|
||||
{"role": "system", "content": "s"},
|
||||
{"role": "user", "content": f"q{i}"},
|
||||
{"role": "assistant", "content": f"a{i}"},
|
||||
]} for i in range(3)]
|
||||
monkeypatch.setattr(trainer_mod, "load_dataset",
|
||||
lambda *a, **k: _FakeDataset(rows), raising=False)
|
||||
monkeypatch.setattr(trainer_mod, "standardize_sharegpt", lambda d: d,
|
||||
raising=False)
|
||||
|
||||
tm = trainer_mod.TrainModel.__new__(trainer_mod.TrainModel)
|
||||
tm.config = {"method": "sft"}
|
||||
tm.chat_tokenizer = _Tokenizer()
|
||||
|
||||
out = tm.process_dataset({"name": "hub/some-dataset"})
|
||||
assert len(out) == 3
|
||||
assert out.rows[0]["text"] == "<system>s<user>q0<assistant>a0"
|
||||
@@ -412,3 +412,88 @@ def test_export_command_bad_target_exits(monkeypatch, tmp_path):
|
||||
quant=None,
|
||||
base_model=None,
|
||||
)
|
||||
|
||||
|
||||
def test_ollama_no_push_skips_push(monkeypatch, tmp_path):
|
||||
"""push=False stops after `ollama create`; nothing is sent to ollama.com."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.setattr(ollama_mod, "_check_ollama_disk", lambda *_a, **_k: None)
|
||||
monkeypatch.setattr(ollama_mod, "ensure_ollama_running", lambda *_a, **_k: None)
|
||||
calls = []
|
||||
|
||||
def _fake_run(cmd, *a, **k):
|
||||
calls.append(cmd)
|
||||
return types.SimpleNamespace(returncode=0, stdout="", stderr="")
|
||||
|
||||
monkeypatch.setattr(ollama_mod.subprocess, "run", _fake_run)
|
||||
ollama_mod.create_and_push_ollama_model(
|
||||
"my-model", "latest", "FROM ./model\n", push=False)
|
||||
assert [c[1] for c in calls] == ["create"]
|
||||
assert not any("push" in c for c in calls)
|
||||
|
||||
|
||||
def _ollama_export_cfg(monkeypatch, tmp_path, push, config_push=None):
|
||||
"""Run `export ollama` with a fake trainer; return the config it received."""
|
||||
import praisonai_train.cli.commands.train as cli
|
||||
|
||||
seen = {}
|
||||
|
||||
class _FakeTrainer:
|
||||
@classmethod
|
||||
def for_export(cls, cfg):
|
||||
obj = cls()
|
||||
obj.config = cfg
|
||||
obj.model = None
|
||||
obj.hf_tokenizer = None
|
||||
return obj
|
||||
|
||||
def load_model(self):
|
||||
return ("MODEL", "TOK")
|
||||
|
||||
def create_and_push_ollama_model(self):
|
||||
seen.update(self.config)
|
||||
|
||||
monkeypatch.setattr(trainer_mod, "TrainModel", _FakeTrainer)
|
||||
config = None
|
||||
if config_push is not None:
|
||||
config = tmp_path / "cfg.yaml"
|
||||
config.write_text(f"ollama_push: {config_push}\n")
|
||||
cli.train_export(
|
||||
target="ollama", model_dir=str(tmp_path), config=str(config) if config else None,
|
||||
ollama="my-model", hf=None, quant=None, base_model="unsloth/gemma-2-2b",
|
||||
mtp_draft=False, push=push,
|
||||
)
|
||||
return seen
|
||||
|
||||
|
||||
def test_export_ollama_no_push_flag_sets_config(monkeypatch, tmp_path):
|
||||
assert _ollama_export_cfg(monkeypatch, tmp_path, push=False)["ollama_push"] is False
|
||||
|
||||
|
||||
def test_export_ollama_explicit_push_overrides_config(monkeypatch, tmp_path):
|
||||
cfg = _ollama_export_cfg(monkeypatch, tmp_path, push=True, config_push="false")
|
||||
assert cfg["ollama_push"] is True
|
||||
|
||||
|
||||
def test_export_ollama_unset_push_keeps_config(monkeypatch, tmp_path):
|
||||
cfg = _ollama_export_cfg(monkeypatch, tmp_path, push=None, config_push="false")
|
||||
assert cfg["ollama_push"] is False # YAML bool preserved, not overridden
|
||||
assert "ollama_push" not in _ollama_export_cfg(monkeypatch, tmp_path, push=None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value, expected", [
|
||||
(False, False), ("false", False), (True, True), ("true", True), (None, True),
|
||||
])
|
||||
def test_llm_trainer_forwards_ollama_push(monkeypatch, value, expected):
|
||||
"""TrainModel forwards `ollama_push` (bool or string) to the helper."""
|
||||
captured = {}
|
||||
monkeypatch.setattr(
|
||||
ollama_mod, "create_and_push_ollama_model",
|
||||
lambda *a, **k: captured.update(k))
|
||||
fake = types.SimpleNamespace(
|
||||
config={"ollama_model": "m", "ollama_push": value},
|
||||
prepare_modelfile_content=lambda: "FROM ./model\n",
|
||||
_flag=trainer_mod.TrainModel._flag,
|
||||
)
|
||||
trainer_mod.TrainModel.create_and_push_ollama_model(fake)
|
||||
assert captured["push"] is expected
|
||||
Generated
+2
-2
@@ -2606,7 +2606,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-train"
|
||||
version = "0.1.20"
|
||||
version = "0.1.21"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -2696,7 +2696,7 @@ provides-extras = ["llm", "vision", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.21"
|
||||
version = "1.7.22"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
# PraisonAI TypeScript Node.js AI Agents Framework
|
||||
|
||||
PraisonAI is a production-ready Multi AI Agents framework, designed to create AI Agents to automate and solve problems ranging from simple tasks to complex challenges. It provides a low-code solution to streamline the building and management of multi-agent LLM systems, emphasising simplicity, customisation, and effective human-agent collaboration.
|
||||
PraisonAI runs AI agents on your own computers, even fully offline, so private data never has to leave the building. It is built for highly regulated teams, and free and open source for everyone.
|
||||
|
||||
This package is the TypeScript and JavaScript SDK of the open-source PraisonAI agent framework. Build single agents, multi-agent workflows, tools and memory in Node.js, on hosted models or on local models such as Ollama. Data stays on your own machines only when the model runs locally too; with a hosted model such as the default OpenAI one, prompts are sent to that provider. The Python framework and the full docs are at [github.com/MervinPraison/PraisonAI](https://github.com/MervinPraison/PraisonAI) and [praison.ai/docs](https://praison.ai/docs).
|
||||
|
||||
## Installation
|
||||
|
||||
|
||||
@@ -46,7 +46,7 @@ This complements `PARITY.md`, which only tracks whether an export exists.
|
||||
|
||||
### `Agent.__init__`
|
||||
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/agent.py:631`
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/agent.py:632`
|
||||
- TypeScript: `src/praisonai-ts/src/agent/simple.ts:197` (ctor `src/praisonai-ts/src/agent/simple.ts:1019`)
|
||||
- Counts: 42 python params: 29 exact, 8 camelCase, 2 alias, 3 flattened, 0 missing; 8 mismatches; 8 waived; 14 TS-only of 60
|
||||
|
||||
@@ -219,7 +219,7 @@ TS-only members: `previousResult`?, `onToken`?, `signal`?, `onEvent`?, `options`
|
||||
|
||||
### `Agent.chat`
|
||||
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3657`
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3734`
|
||||
- TypeScript: `src/praisonai-ts/src/agent/simple.ts:3592`
|
||||
- Counts: 17 python params: 7 exact, 9 camelCase, 1 alias, 0 flattened, 0 missing; 0 mismatches; 0 waived; 2 TS-only of 19
|
||||
|
||||
@@ -316,7 +316,7 @@ TS-only members: `config`
|
||||
### `Session.__init__`
|
||||
|
||||
- Python: `src/praisonai-agents/praisonaiagents/session/api.py:67`
|
||||
- TypeScript: `src/praisonai-ts/src/session/session.ts:29` (ctor `src/praisonai-ts/src/session/session.ts:132`)
|
||||
- TypeScript: `src/praisonai-ts/src/session/session.ts:29` (ctor `src/praisonai-ts/src/session/session.ts:133`)
|
||||
- Counts: 7 python params: 1 exact, 6 camelCase, 0 alias, 0 flattened, 0 missing; 1 mismatches; 1 waived; 5 TS-only of 12
|
||||
|
||||
| Python param | Kind | Py default | Py type | Match | TS name | TS default | TS type | Status |
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"name": "praisonai",
|
||||
"version": "1.7.4",
|
||||
"description": "PraisonAI TypeScript AI Agents Framework - Node.js, npm, and Javascript AI Agents Framework",
|
||||
"description": "Open-source TypeScript and JavaScript AI agent framework for Node.js, on hosted or local models. The JavaScript SDK of PraisonAI.",
|
||||
"main": "dist/index.js",
|
||||
"module": "dist/esm/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
|
||||
@@ -1185,7 +1185,7 @@
|
||||
}
|
||||
],
|
||||
"python": {
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/agent.py:631",
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/agent.py:632",
|
||||
"resolved_class": "Agent"
|
||||
},
|
||||
"ts_only": [
|
||||
@@ -1797,7 +1797,7 @@
|
||||
}
|
||||
],
|
||||
"python": {
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3657",
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3734",
|
||||
"resolved_class": "ChatMixin"
|
||||
},
|
||||
"ts_only": [
|
||||
@@ -5142,7 +5142,7 @@
|
||||
],
|
||||
"ts_only_required": [],
|
||||
"typescript": {
|
||||
"ctor_location": "src/praisonai-ts/src/session/session.ts:132",
|
||||
"ctor_location": "src/praisonai-ts/src/session/session.ts:133",
|
||||
"location": "src/praisonai-ts/src/session/session.ts:29",
|
||||
"options_parameters": [
|
||||
"config"
|
||||
|
||||
Vendored
+20
-5
@@ -72,7 +72,7 @@ export class MemoryCache extends BaseCache {
|
||||
if (this.maxSize && this.cache.size >= this.maxSize && !this.cache.has(key)) {
|
||||
// Remove oldest entry
|
||||
const firstKey = this.cache.keys().next().value;
|
||||
if (firstKey) this.cache.delete(firstKey);
|
||||
if (firstKey !== undefined) this.cache.delete(firstKey);
|
||||
}
|
||||
|
||||
const effectiveTtl = ttl ?? this.ttl;
|
||||
@@ -137,12 +137,15 @@ export class MemoryCache extends BaseCache {
|
||||
*/
|
||||
export class FileCache extends BaseCache {
|
||||
private cacheDir: string;
|
||||
private ttl?: number;
|
||||
private fs: any;
|
||||
private path: any;
|
||||
private listQueues: Map<string, Promise<void>> = new Map();
|
||||
|
||||
constructor(config: CacheConfig & { cacheDir?: string } = {}) {
|
||||
super(config);
|
||||
this.cacheDir = config.cacheDir || '.cache';
|
||||
this.ttl = config.ttl;
|
||||
}
|
||||
|
||||
private async ensureDir(): Promise<void> {
|
||||
@@ -177,10 +180,11 @@ export class FileCache extends BaseCache {
|
||||
|
||||
async set<T = any>(key: string, value: T, ttl?: number): Promise<void> {
|
||||
await this.ensureDir();
|
||||
const effectiveTtl = ttl ?? this.ttl;
|
||||
const entry: CacheEntry<T> = {
|
||||
value,
|
||||
createdAt: Date.now(),
|
||||
expiresAt: ttl ? Date.now() + ttl : undefined
|
||||
expiresAt: effectiveTtl ? Date.now() + effectiveTtl : undefined
|
||||
};
|
||||
await this.fs.writeFile(this.getFilePath(key), JSON.stringify(entry));
|
||||
}
|
||||
@@ -227,9 +231,20 @@ export class FileCache extends BaseCache {
|
||||
}
|
||||
|
||||
async listPush<T = any>(key: string, value: T): Promise<void> {
|
||||
const list = (await this.get<T[]>(`list_${key}`)) || [];
|
||||
list.push(value);
|
||||
await this.set(`list_${key}`, list);
|
||||
// Serialize the full read/modify/write operation for this instance and key.
|
||||
// A failed append must not block subsequent appends.
|
||||
const previous = this.listQueues.get(key) ?? Promise.resolve();
|
||||
const next = previous.catch(() => {}).then(async () => {
|
||||
const list = (await this.get<T[]>(`list_${key}`)) || [];
|
||||
list.push(value);
|
||||
await this.set(`list_${key}`, list);
|
||||
});
|
||||
this.listQueues.set(key, next);
|
||||
try {
|
||||
await next;
|
||||
} finally {
|
||||
if (this.listQueues.get(key) === next) this.listQueues.delete(key);
|
||||
}
|
||||
}
|
||||
|
||||
async listLength(key: string): Promise<number> {
|
||||
|
||||
@@ -231,12 +231,17 @@ export class HooksManager {
|
||||
return promise;
|
||||
}
|
||||
|
||||
return Promise.race([
|
||||
promise,
|
||||
new Promise<T>((_, reject) =>
|
||||
setTimeout(() => reject(new Error(`Hook timeout after ${timeout}ms`)), timeout)
|
||||
),
|
||||
]);
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
try {
|
||||
return await Promise.race([
|
||||
promise,
|
||||
new Promise<T>((_, reject) => {
|
||||
timer = setTimeout(() => reject(new Error(`Hook timeout after ${timeout}ms`)), timeout);
|
||||
}),
|
||||
]);
|
||||
} finally {
|
||||
if (timer !== undefined) clearTimeout(timer);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -251,7 +256,7 @@ export class HooksManager {
|
||||
* Get all registered hooks for an event
|
||||
*/
|
||||
getHooks(event: HookEvent): readonly HookConfig[] {
|
||||
return Object.freeze(this.hooks.get(event) ?? []);
|
||||
return Object.freeze([...(this.hooks.get(event) ?? [])]);
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -82,16 +82,20 @@ export class Chunking {
|
||||
* Chunk by sentences
|
||||
*/
|
||||
chunkBySentence(text: string): Chunk[] {
|
||||
const sentences = text.match(/[^.!?]+[.!?]+/g) || [text];
|
||||
return sentences.map((content, index) => {
|
||||
const startOffset = text.indexOf(content);
|
||||
return {
|
||||
content: content.trim(),
|
||||
index,
|
||||
const chunks: Chunk[] = [];
|
||||
// Include the unterminated tail and punctuation-only spans as well.
|
||||
for (const match of text.matchAll(/[^.!?]+[.!?]*|[.!?]+/g)) {
|
||||
const content = match[0].trim();
|
||||
if (!content) continue;
|
||||
const startOffset = match.index + match[0].indexOf(content);
|
||||
chunks.push({
|
||||
content,
|
||||
index: chunks.length,
|
||||
startOffset,
|
||||
endOffset: startOffset + content.length
|
||||
};
|
||||
});
|
||||
});
|
||||
}
|
||||
return chunks;
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -60,7 +60,23 @@ export class GraphStore {
|
||||
* Add an edge to the graph
|
||||
*/
|
||||
addEdge(edge: GraphEdge): void {
|
||||
const previous = this.edges.get(edge.id);
|
||||
this.edges.set(edge.id, edge);
|
||||
|
||||
if (previous && (previous.source !== edge.source || previous.target !== edge.target)) {
|
||||
// Another edge may still connect the old endpoints.
|
||||
let hasParallelEdge = false;
|
||||
for (const candidate of this.edges.values()) {
|
||||
if (candidate.source === previous.source && candidate.target === previous.target) {
|
||||
hasParallelEdge = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!hasParallelEdge) {
|
||||
this.adjacencyList.get(previous.source)?.delete(previous.target);
|
||||
this.reverseAdjacencyList.get(previous.target)?.delete(previous.source);
|
||||
}
|
||||
}
|
||||
|
||||
// Update adjacency lists
|
||||
if (!this.adjacencyList.has(edge.source)) {
|
||||
|
||||
@@ -70,22 +70,24 @@ export class WebSocketTransport implements MCPTransport {
|
||||
}
|
||||
const WebSocket = wsModule.default;
|
||||
|
||||
this.ws = new WebSocket(this.url);
|
||||
const socket = new WebSocket(this.url);
|
||||
this.ws = socket;
|
||||
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
this.ws.on('open', () => {
|
||||
socket.on('open', () => {
|
||||
this.connected = true;
|
||||
resolve();
|
||||
});
|
||||
this.ws.on('error', (err: Error) => reject(err));
|
||||
socket.on('error', (err: Error) => reject(err));
|
||||
});
|
||||
|
||||
this.ws.on('message', (data: Buffer | string) => {
|
||||
socket.on('message', (data: Buffer | string) => {
|
||||
if (this.ws !== socket) return;
|
||||
try {
|
||||
const message = JSON.parse(data.toString()) as TransportMessage;
|
||||
|
||||
// Handle response
|
||||
if (message.id && this.pendingRequests.has(String(message.id))) {
|
||||
if (message.id !== undefined && this.pendingRequests.has(String(message.id))) {
|
||||
const { resolve, reject } = this.pendingRequests.get(String(message.id))!;
|
||||
this.pendingRequests.delete(String(message.id));
|
||||
if (message.error) {
|
||||
@@ -102,16 +104,29 @@ export class WebSocketTransport implements MCPTransport {
|
||||
}
|
||||
});
|
||||
|
||||
this.ws.on('error', (err: Error) => this.errorHandler?.(err));
|
||||
this.ws.on('close', () => { this.connected = false; });
|
||||
socket.on('error', (err: Error) => {
|
||||
if (this.ws === socket) this.errorHandler?.(err);
|
||||
});
|
||||
socket.on('close', () => {
|
||||
if (this.ws !== socket) return;
|
||||
this.connected = false;
|
||||
this.ws = null;
|
||||
this.rejectPendingRequests();
|
||||
});
|
||||
}
|
||||
|
||||
async disconnect(): Promise<void> {
|
||||
if (this.ws) {
|
||||
this.ws.close();
|
||||
this.ws = null;
|
||||
this.connected = false;
|
||||
}
|
||||
const socket = this.ws;
|
||||
this.ws = null;
|
||||
this.connected = false;
|
||||
this.rejectPendingRequests();
|
||||
socket?.close();
|
||||
}
|
||||
|
||||
private rejectPendingRequests(): void {
|
||||
const error = new Error('WebSocket connection closed');
|
||||
for (const { reject } of this.pendingRequests.values()) reject(error);
|
||||
this.pendingRequests.clear();
|
||||
}
|
||||
|
||||
async send(message: TransportMessage): Promise<TransportMessage> {
|
||||
|
||||
@@ -28,6 +28,7 @@ export class FileMemory {
|
||||
private initialized: boolean = false;
|
||||
private initPromise?: Promise<void>;
|
||||
private writeQueue: Promise<void> = Promise.resolve();
|
||||
private appendNeedsSeparator: boolean = false;
|
||||
|
||||
constructor(config: FileMemoryConfig) {
|
||||
this.filePath = config.filePath;
|
||||
@@ -67,6 +68,8 @@ export class FileMemory {
|
||||
try {
|
||||
const fs = await import('fs/promises');
|
||||
const content = await fs.readFile(this.filePath, 'utf-8').catch(() => '');
|
||||
// A previous instance may have left a record without its final newline.
|
||||
this.appendNeedsSeparator = content.length > 0 && !content.endsWith('\n');
|
||||
|
||||
if (content) {
|
||||
const lines = content.trim().split('\n').filter(Boolean);
|
||||
@@ -155,6 +158,7 @@ export class FileMemory {
|
||||
*/
|
||||
async getRecent(count: number): Promise<MemoryEntry[]> {
|
||||
const all = await this.getAll();
|
||||
if (count === 0) return [];
|
||||
return all.slice(-count);
|
||||
}
|
||||
|
||||
@@ -265,7 +269,7 @@ export class FileMemory {
|
||||
* Build context string from recent memory
|
||||
*/
|
||||
async buildContext(count?: number): Promise<string> {
|
||||
const entries = count ? await this.getRecent(count) : await this.getAll();
|
||||
const entries = count !== undefined ? await this.getRecent(count) : await this.getAll();
|
||||
return entries
|
||||
.map(e => `${e.role}: ${e.content}`)
|
||||
.join('\n');
|
||||
@@ -273,11 +277,17 @@ export class FileMemory {
|
||||
|
||||
private async appendEntry(entry: FileMemoryEntry): Promise<void> {
|
||||
// Queue writes to ensure order
|
||||
this.writeQueue = this.writeQueue.then(async () => {
|
||||
const write = this.writeQueue.then(async () => {
|
||||
const fs = await import('fs/promises');
|
||||
await fs.appendFile(this.filePath, JSON.stringify(entry) + '\n');
|
||||
// A failed append may have left an unfinished JSONL line.
|
||||
const separator = this.appendNeedsSeparator ? '\n' : '';
|
||||
this.appendNeedsSeparator = true;
|
||||
await fs.appendFile(this.filePath, separator + JSON.stringify(entry) + '\n');
|
||||
this.appendNeedsSeparator = false;
|
||||
});
|
||||
await this.writeQueue;
|
||||
// Report this write's error to its caller without blocking later writes.
|
||||
this.writeQueue = write.catch(() => {});
|
||||
await write;
|
||||
}
|
||||
|
||||
private textSearch(query: string, limit: number): SearchResult[] {
|
||||
|
||||
@@ -87,6 +87,7 @@ export class Memory {
|
||||
* Get recent entries
|
||||
*/
|
||||
getRecent(count: number): MemoryEntry[] {
|
||||
if (count === 0) return [];
|
||||
return this.getAll().slice(-count);
|
||||
}
|
||||
|
||||
@@ -179,7 +180,7 @@ export class Memory {
|
||||
* Build context string from recent memory
|
||||
*/
|
||||
buildContext(count?: number): string {
|
||||
const entries = count ? this.getRecent(count) : this.getAll();
|
||||
const entries = count !== undefined ? this.getRecent(count) : this.getAll();
|
||||
return entries
|
||||
.map(e => `${e.role}: ${e.content}`)
|
||||
.join('\n');
|
||||
|
||||
Loaded 100 of 135 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user