This is page 13 of 14. Use http://codebase.md/stanfordnlp/dspy?page={x} to view the full context.
# Directory Structure
```
├── .github
│ ├── .internal_dspyai
│ │ ├── internals
│ │ │ ├── build-and-release.md
│ │ │ └── release-checklist.md
│ │ └── pyproject.toml
│ ├── .tmp
│ │ └── .generated-actions
│ │ └── run-pypi-publish-in-docker-container
│ │ └── action.yml
│ ├── ISSUE_TEMPLATE
│ │ ├── bug_report.yml
│ │ └── feature_request.yml
│ ├── PULL_REQUEST_TEMPLATE
│ │ └── pull_request_template.md
│ ├── workflow_scripts
│ │ └── install_testpypi_pkg.sh
│ └── workflows
│ ├── build_and_release.yml
│ ├── build_utils
│ │ └── test_version.py
│ ├── docs-push.yml
│ ├── precommits_check.yml
│ └── run_tests.yml
├── .gitignore
├── .pre-commit-config.yaml
├── CONTRIBUTING.md
├── docs
│ ├── .gitignore
│ ├── docs
│ │ ├── api
│ │ │ ├── adapters
│ │ │ │ ├── Adapter.md
│ │ │ │ ├── ChatAdapter.md
│ │ │ │ ├── JSONAdapter.md
│ │ │ │ └── TwoStepAdapter.md
│ │ │ ├── evaluation
│ │ │ │ ├── answer_exact_match.md
│ │ │ │ ├── answer_passage_match.md
│ │ │ │ ├── CompleteAndGrounded.md
│ │ │ │ ├── Evaluate.md
│ │ │ │ ├── EvaluationResult.md
│ │ │ │ └── SemanticF1.md
│ │ │ ├── experimental
│ │ │ │ ├── Citations.md
│ │ │ │ └── Document.md
│ │ │ ├── index.md
│ │ │ ├── models
│ │ │ │ ├── Embedder.md
│ │ │ │ └── LM.md
│ │ │ ├── modules
│ │ │ │ ├── BestOfN.md
│ │ │ │ ├── ChainOfThought.md
│ │ │ │ ├── CodeAct.md
│ │ │ │ ├── Module.md
│ │ │ │ ├── MultiChainComparison.md
│ │ │ │ ├── Parallel.md
│ │ │ │ ├── Predict.md
│ │ │ │ ├── ProgramOfThought.md
│ │ │ │ ├── ReAct.md
│ │ │ │ └── Refine.md
│ │ │ ├── optimizers
│ │ │ │ ├── BetterTogether.md
│ │ │ │ ├── BootstrapFewShot.md
│ │ │ │ ├── BootstrapFewShotWithRandomSearch.md
│ │ │ │ ├── BootstrapFinetune.md
│ │ │ │ ├── BootstrapRS.md
│ │ │ │ ├── COPRO.md
│ │ │ │ ├── Ensemble.md
│ │ │ │ ├── GEPA
│ │ │ │ │ ├── GEPA_Advanced.md
│ │ │ │ │ └── overview.md
│ │ │ │ ├── InferRules.md
│ │ │ │ ├── KNN.md
│ │ │ │ ├── KNNFewShot.md
│ │ │ │ ├── LabeledFewShot.md
│ │ │ │ ├── MIPROv2.md
│ │ │ │ └── SIMBA.md
│ │ │ ├── primitives
│ │ │ │ ├── Audio.md
│ │ │ │ ├── Code.md
│ │ │ │ ├── Example.md
│ │ │ │ ├── History.md
│ │ │ │ ├── Image.md
│ │ │ │ ├── Prediction.md
│ │ │ │ ├── Tool.md
│ │ │ │ └── ToolCalls.md
│ │ │ ├── signatures
│ │ │ │ ├── InputField.md
│ │ │ │ ├── OutputField.md
│ │ │ │ └── Signature.md
│ │ │ ├── tools
│ │ │ │ ├── ColBERTv2.md
│ │ │ │ ├── Embeddings.md
│ │ │ │ └── PythonInterpreter.md
│ │ │ └── utils
│ │ │ ├── asyncify.md
│ │ │ ├── configure_cache.md
│ │ │ ├── disable_litellm_logging.md
│ │ │ ├── disable_logging.md
│ │ │ ├── enable_litellm_logging.md
│ │ │ ├── enable_logging.md
│ │ │ ├── inspect_history.md
│ │ │ ├── load.md
│ │ │ ├── StatusMessage.md
│ │ │ ├── StatusMessageProvider.md
│ │ │ ├── streamify.md
│ │ │ └── StreamListener.md
│ │ ├── cheatsheet.md
│ │ ├── community
│ │ │ ├── community-resources.md
│ │ │ ├── how-to-contribute.md
│ │ │ └── use-cases.md
│ │ ├── deep-dive
│ │ │ └── data-handling
│ │ │ ├── built-in-datasets.md
│ │ │ ├── examples.md
│ │ │ ├── img
│ │ │ │ └── data-loading.png
│ │ │ └── loading-custom-data.md
│ │ ├── faqs.md
│ │ ├── index.md
│ │ ├── js
│ │ │ └── runllm-widget.js
│ │ ├── learn
│ │ │ ├── evaluation
│ │ │ │ ├── data.md
│ │ │ │ ├── metrics.md
│ │ │ │ └── overview.md
│ │ │ ├── figures
│ │ │ │ ├── native_tool_call.png
│ │ │ │ └── teleprompter-classes.png
│ │ │ ├── index.md
│ │ │ ├── optimization
│ │ │ │ ├── optimizers.md
│ │ │ │ └── overview.md
│ │ │ └── programming
│ │ │ ├── 7-assertions.md
│ │ │ ├── adapters.md
│ │ │ ├── language_models.md
│ │ │ ├── mcp.md
│ │ │ ├── modules.md
│ │ │ ├── overview.md
│ │ │ ├── signatures.md
│ │ │ └── tools.md
│ │ ├── production
│ │ │ └── index.md
│ │ ├── roadmap.md
│ │ ├── static
│ │ │ ├── .nojekyll
│ │ │ └── img
│ │ │ ├── dspy_logo.png
│ │ │ ├── logo.png
│ │ │ ├── mlflow-tracing-rag.png
│ │ │ ├── modular.png
│ │ │ ├── optimize.png
│ │ │ ├── undraw_docusaurus_mountain.svg
│ │ │ ├── undraw_docusaurus_react.svg
│ │ │ ├── undraw_docusaurus_tree.svg
│ │ │ └── universal_compatibility.png
│ │ ├── stylesheets
│ │ │ └── extra.css
│ │ └── tutorials
│ │ ├── agents
│ │ │ ├── index.ipynb
│ │ │ └── mlflow-tracing-agent.png
│ │ ├── ai_text_game
│ │ │ └── index.md
│ │ ├── async
│ │ │ └── index.md
│ │ ├── audio
│ │ │ └── index.ipynb
│ │ ├── build_ai_program
│ │ │ └── index.md
│ │ ├── cache
│ │ │ └── index.md
│ │ ├── classification
│ │ │ └── index.md
│ │ ├── classification_finetuning
│ │ │ ├── index.ipynb
│ │ │ └── mlflow-tracing-classification.png
│ │ ├── conversation_history
│ │ │ └── index.md
│ │ ├── core_development
│ │ │ └── index.md
│ │ ├── custom_module
│ │ │ ├── index.ipynb
│ │ │ └── mlflow-tracing-custom-module.png
│ │ ├── customer_service_agent
│ │ │ ├── index.ipynb
│ │ │ └── mlflow-tracing-customer-service-agent.png
│ │ ├── deployment
│ │ │ ├── dspy_mlflow_ui.png
│ │ │ └── index.md
│ │ ├── email_extraction
│ │ │ ├── index.md
│ │ │ └── mlflow-tracing-email-extraction.png
│ │ ├── entity_extraction
│ │ │ ├── index.ipynb
│ │ │ └── mlflow-tracing-entity-extraction.png
│ │ ├── games
│ │ │ ├── index.ipynb
│ │ │ └── mlflow-tracing-agent.png
│ │ ├── gepa_ai_program
│ │ │ └── index.md
│ │ ├── gepa_aime
│ │ │ ├── index.ipynb
│ │ │ ├── mlflow-tracing-gepa-aime.png
│ │ │ └── mlflow-tracking-gepa-aime-optimization.png
│ │ ├── gepa_facilitysupportanalyzer
│ │ │ ├── index.ipynb
│ │ │ ├── mlflow-tracing-gepa-support.png
│ │ │ └── mlflow-tracking-gepa-support-optimization.png
│ │ ├── gepa_papillon
│ │ │ ├── index.ipynb
│ │ │ ├── mlflow-tracing-gepa-papilon.png
│ │ │ └── mlflow-tracking-gepa-papilon-optimization.png
│ │ ├── image_generation_prompting
│ │ │ └── index.ipynb
│ │ ├── index.md
│ │ ├── llms_txt_generation
│ │ │ └── index.md
│ │ ├── math
│ │ │ ├── index.ipynb
│ │ │ └── mlflow-tracing-math.png
│ │ ├── mcp
│ │ │ └── index.md
│ │ ├── mem0_react_agent
│ │ │ └── index.md
│ │ ├── multihop_search
│ │ │ ├── index.ipynb
│ │ │ └── mlflow-tracing-multi-hop.png
│ │ ├── observability
│ │ │ ├── index.md
│ │ │ ├── mlflow_trace_ui_navigation.gif
│ │ │ ├── mlflow_trace_ui.png
│ │ │ └── mlflow_trace_view.png
│ │ ├── optimize_ai_program
│ │ │ └── index.md
│ │ ├── optimizer_tracking
│ │ │ ├── child_run.png
│ │ │ ├── experiment.png
│ │ │ ├── index.md
│ │ │ └── parent_run.png
│ │ ├── output_refinement
│ │ │ └── best-of-n-and-refine.md
│ │ ├── papillon
│ │ │ └── index.md
│ │ ├── program_of_thought
│ │ │ └── index.ipynb
│ │ ├── rag
│ │ │ ├── index.ipynb
│ │ │ └── mlflow-tracing-rag.png
│ │ ├── real_world_examples
│ │ │ └── index.md
│ │ ├── rl_ai_program
│ │ │ └── index.md
│ │ ├── rl_multihop
│ │ │ └── index.ipynb
│ │ ├── rl_papillon
│ │ │ └── index.ipynb
│ │ ├── sample_code_generation
│ │ │ └── index.md
│ │ ├── saving
│ │ │ └── index.md
│ │ ├── streaming
│ │ │ └── index.md
│ │ ├── tool_use
│ │ │ └── index.ipynb
│ │ └── yahoo_finance_react
│ │ └── index.md
│ ├── mkdocs.yml
│ ├── overrides
│ │ ├── home.html
│ │ ├── main.html
│ │ └── partials
│ │ └── tabs.html
│ ├── Pipfile
│ ├── Pipfile.lock
│ ├── README.md
│ ├── requirements.txt
│ ├── scripts
│ │ ├── generate_api_docs.py
│ │ └── generate_api_summary.py
│ └── vercel.json
├── dspy
│ ├── __init__.py
│ ├── __metadata__.py
│ ├── adapters
│ │ ├── __init__.py
│ │ ├── baml_adapter.py
│ │ ├── base.py
│ │ ├── chat_adapter.py
│ │ ├── json_adapter.py
│ │ ├── two_step_adapter.py
│ │ ├── types
│ │ │ ├── __init__.py
│ │ │ ├── audio.py
│ │ │ ├── base_type.py
│ │ │ ├── citation.py
│ │ │ ├── code.py
│ │ │ ├── document.py
│ │ │ ├── history.py
│ │ │ ├── image.py
│ │ │ └── tool.py
│ │ ├── utils.py
│ │ └── xml_adapter.py
│ ├── clients
│ │ ├── __init__.py
│ │ ├── base_lm.py
│ │ ├── cache.py
│ │ ├── databricks.py
│ │ ├── embedding.py
│ │ ├── lm_local_arbor.py
│ │ ├── lm_local.py
│ │ ├── lm.py
│ │ ├── openai.py
│ │ ├── provider.py
│ │ └── utils_finetune.py
│ ├── datasets
│ │ ├── __init__.py
│ │ ├── alfworld
│ │ │ ├── __init__.py
│ │ │ ├── alfworld.py
│ │ │ └── base_config.yml
│ │ ├── colors.py
│ │ ├── dataloader.py
│ │ ├── dataset.py
│ │ ├── gsm8k.py
│ │ ├── hotpotqa.py
│ │ └── math.py
│ ├── dsp
│ │ ├── __init__.py
│ │ ├── colbertv2.py
│ │ └── utils
│ │ ├── __init__.py
│ │ ├── dpr.py
│ │ ├── settings.py
│ │ └── utils.py
│ ├── evaluate
│ │ ├── __init__.py
│ │ ├── auto_evaluation.py
│ │ ├── evaluate.py
│ │ └── metrics.py
│ ├── experimental
│ │ └── __init__.py
│ ├── predict
│ │ ├── __init__.py
│ │ ├── aggregation.py
│ │ ├── avatar
│ │ │ ├── __init__.py
│ │ │ ├── avatar.py
│ │ │ ├── models.py
│ │ │ └── signatures.py
│ │ ├── best_of_n.py
│ │ ├── chain_of_thought.py
│ │ ├── code_act.py
│ │ ├── knn.py
│ │ ├── multi_chain_comparison.py
│ │ ├── parallel.py
│ │ ├── parameter.py
│ │ ├── predict.py
│ │ ├── program_of_thought.py
│ │ ├── react.py
│ │ ├── refine.py
│ │ └── retry.py
│ ├── primitives
│ │ ├── __init__.py
│ │ ├── base_module.py
│ │ ├── example.py
│ │ ├── module.py
│ │ ├── prediction.py
│ │ ├── python_interpreter.py
│ │ └── runner.js
│ ├── propose
│ │ ├── __init__.py
│ │ ├── dataset_summary_generator.py
│ │ ├── grounded_proposer.py
│ │ ├── propose_base.py
│ │ └── utils.py
│ ├── retrievers
│ │ ├── __init__.py
│ │ ├── databricks_rm.py
│ │ ├── embeddings.py
│ │ ├── retrieve.py
│ │ └── weaviate_rm.py
│ ├── signatures
│ │ ├── __init__.py
│ │ ├── field.py
│ │ ├── signature.py
│ │ └── utils.py
│ ├── streaming
│ │ ├── __init__.py
│ │ ├── messages.py
│ │ ├── streamify.py
│ │ └── streaming_listener.py
│ ├── teleprompt
│ │ ├── __init__.py
│ │ ├── avatar_optimizer.py
│ │ ├── bettertogether.py
│ │ ├── bootstrap_finetune.py
│ │ ├── bootstrap_trace.py
│ │ ├── bootstrap.py
│ │ ├── copro_optimizer.py
│ │ ├── ensemble.py
│ │ ├── gepa
│ │ │ ├── __init__.py
│ │ │ ├── gepa_utils.py
│ │ │ ├── gepa.py
│ │ │ └── instruction_proposal.py
│ │ ├── grpo.py
│ │ ├── infer_rules.py
│ │ ├── knn_fewshot.py
│ │ ├── mipro_optimizer_v2.py
│ │ ├── random_search.py
│ │ ├── signature_opt.py
│ │ ├── simba_utils.py
│ │ ├── simba.py
│ │ ├── teleprompt_optuna.py
│ │ ├── teleprompt.py
│ │ ├── utils.py
│ │ └── vanilla.py
│ └── utils
│ ├── __init__.py
│ ├── annotation.py
│ ├── asyncify.py
│ ├── caching.py
│ ├── callback.py
│ ├── dummies.py
│ ├── exceptions.py
│ ├── hasher.py
│ ├── inspect_history.py
│ ├── langchain_tool.py
│ ├── logging_utils.py
│ ├── mcp.py
│ ├── parallelizer.py
│ ├── saving.py
│ ├── syncify.py
│ ├── unbatchify.py
│ └── usage_tracker.py
├── LICENSE
├── pyproject.toml
├── README.md
├── tests
│ ├── __init__.py
│ ├── adapters
│ │ ├── test_adapter_utils.py
│ │ ├── test_baml_adapter.py
│ │ ├── test_base_type.py
│ │ ├── test_chat_adapter.py
│ │ ├── test_citation.py
│ │ ├── test_code.py
│ │ ├── test_document.py
│ │ ├── test_json_adapter.py
│ │ ├── test_tool.py
│ │ ├── test_two_step_adapter.py
│ │ └── test_xml_adapter.py
│ ├── callback
│ │ └── test_callback.py
│ ├── clients
│ │ ├── test_cache.py
│ │ ├── test_databricks.py
│ │ ├── test_embedding.py
│ │ ├── test_inspect_global_history.py
│ │ └── test_lm.py
│ ├── conftest.py
│ ├── datasets
│ │ └── test_dataset.py
│ ├── docs
│ │ └── test_mkdocs_links.py
│ ├── evaluate
│ │ ├── test_evaluate.py
│ │ └── test_metrics.py
│ ├── examples
│ │ └── test_baleen.py
│ ├── metadata
│ │ └── test_metadata.py
│ ├── predict
│ │ ├── test_aggregation.py
│ │ ├── test_best_of_n.py
│ │ ├── test_chain_of_thought.py
│ │ ├── test_code_act.py
│ │ ├── test_knn.py
│ │ ├── test_multi_chain_comparison.py
│ │ ├── test_parallel.py
│ │ ├── test_predict.py
│ │ ├── test_program_of_thought.py
│ │ ├── test_react.py
│ │ ├── test_refine.py
│ │ └── test_retry.py
│ ├── primitives
│ │ ├── resources
│ │ │ └── saved_program.json
│ │ ├── test_base_module.py
│ │ ├── test_example.py
│ │ ├── test_module.py
│ │ └── test_python_interpreter.py
│ ├── propose
│ │ └── test_grounded_proposer.py
│ ├── README.md
│ ├── reliability
│ │ ├── __init__.py
│ │ ├── complex_types
│ │ │ └── generated
│ │ │ ├── test_many_types_1
│ │ │ │ ├── inputs
│ │ │ │ │ ├── input1.json
│ │ │ │ │ └── input2.json
│ │ │ │ ├── program.py
│ │ │ │ └── schema.json
│ │ │ ├── test_nesting_1
│ │ │ │ ├── inputs
│ │ │ │ │ ├── input1.json
│ │ │ │ │ └── input2.json
│ │ │ │ ├── program.py
│ │ │ │ └── schema.json
│ │ │ └── test_nesting_2
│ │ │ ├── inputs
│ │ │ │ └── input1.json
│ │ │ ├── program.py
│ │ │ └── schema.json
│ │ ├── conftest.py
│ │ ├── generate
│ │ │ ├── __init__.py
│ │ │ ├── __main__.py
│ │ │ └── utils.py
│ │ ├── input_formats
│ │ │ └── generated
│ │ │ └── test_markdown_1
│ │ │ ├── inputs
│ │ │ │ ├── input1.json
│ │ │ │ └── input2.json
│ │ │ ├── program.py
│ │ │ └── schema.json
│ │ ├── README.md
│ │ ├── reliability_conf.yaml
│ │ ├── test_generated.py
│ │ ├── test_pydantic_models.py
│ │ └── utils.py
│ ├── retrievers
│ │ └── test_embeddings.py
│ ├── signatures
│ │ ├── test_adapter_image.py
│ │ ├── test_custom_types.py
│ │ └── test_signature.py
│ ├── streaming
│ │ └── test_streaming.py
│ ├── teleprompt
│ │ ├── gepa_dummy_lm_custom_component_selector_custom_instruction_proposer.json
│ │ ├── gepa_dummy_lm.json
│ │ ├── test_bootstrap_finetune.py
│ │ ├── test_bootstrap_trace.py
│ │ ├── test_bootstrap.py
│ │ ├── test_copro_optimizer.py
│ │ ├── test_ensemble.py
│ │ ├── test_finetune.py
│ │ ├── test_gepa_instruction_proposer.py
│ │ ├── test_gepa.py
│ │ ├── test_grpo.py
│ │ ├── test_knn_fewshot.py
│ │ ├── test_random_search.py
│ │ ├── test_teleprompt.py
│ │ └── test_utils.py
│ ├── test_utils
│ │ ├── __init__.py
│ │ └── server
│ │ ├── __init__.py
│ │ ├── litellm_server_config.yaml
│ │ └── litellm_server.py
│ └── utils
│ ├── __init__.py
│ ├── resources
│ │ └── mcp_server.py
│ ├── test_annotation.py
│ ├── test_asyncify.py
│ ├── test_exceptions.py
│ ├── test_langchain_tool.py
│ ├── test_mcp.py
│ ├── test_parallelizer.py
│ ├── test_saving.py
│ ├── test_settings.py
│ ├── test_syncify.py
│ ├── test_unbatchify.py
│ └── test_usage_tracker.py
└── uv.lock
```
# Files
--------------------------------------------------------------------------------
/dspy/teleprompt/grpo.py:
--------------------------------------------------------------------------------
```python
import logging
import random
from collections import Counter
from typing import Any, Callable, Literal
from dspy.adapters.base import Adapter
from dspy.adapters.chat_adapter import ChatAdapter
from dspy.clients.lm import LM
from dspy.clients.utils_finetune import GRPOGroup, MultiGPUConfig, TrainDataFormat
from dspy.dsp.utils.settings import settings
from dspy.evaluate.evaluate import Evaluate
from dspy.primitives.example import Example
from dspy.primitives.module import Module
from dspy.teleprompt.bootstrap_finetune import (
FinetuneTeleprompter,
all_predictors_have_lms,
assert_structural_equivalency,
)
from dspy.teleprompt.bootstrap_trace import FailedPrediction, bootstrap_trace_data
logger = logging.getLogger(__name__)
class GRPO(FinetuneTeleprompter):
def __init__(
self,
metric: Callable | None = None,
multitask: bool = True,
train_kwargs: dict[str, Any] | dict[LM, dict[str, Any]] | None = None,
adapter: Adapter | dict[LM, Adapter] | None = None,
exclude_demos: bool = False,
num_threads: int = 6,
num_train_steps: int = 100,
seed: int = 0,
num_dspy_examples_per_grpo_step: int = 1,
num_rollouts_per_grpo_step: int = 1,
use_train_as_val: bool = False,
num_steps_for_val: int = 5,
report_train_scores: bool = False,
failure_score: float = 0,
format_failure_score: float = -1,
variably_invoked_predictor_grouping_mode: Literal["truncate"] | Literal["fill"] | Literal["ragged"] = "truncate",
variably_invoked_predictor_fill_strategy: Literal["randint"] | Literal["max"] | None = None,
gpu_config: MultiGPUConfig = MultiGPUConfig(num_inference_gpus=1, num_training_gpus=1),
):
super().__init__(train_kwargs=train_kwargs)
self.metric = metric
self.multitask = multitask
self.adapter: dict[LM, Adapter] = self.convert_to_lm_dict(adapter)
self.exclude_demos = exclude_demos
self.num_threads = num_threads
self.num_train_steps = num_train_steps
self.rng = random.Random(seed)
self.num_dspy_examples_per_grpo_step = num_dspy_examples_per_grpo_step
self.num_rollouts_per_grpo_step = num_rollouts_per_grpo_step
self.use_train_as_val = use_train_as_val
self.num_steps_for_val = num_steps_for_val
self.report_train_scores = report_train_scores
self.failure_score = failure_score
self.format_failure_score = format_failure_score
self.gpu_config = gpu_config
assert failure_score > format_failure_score, "failure_score must be greater than format_failure_score since the range [format_failure_score, failure_score] is used to provide dspy formatting rewards"
if self.use_train_as_val:
assert report_train_scores, "If use_train_as_val is True, report_train_scores must be True."
assert exclude_demos, "exclude_demos==False is not supported yet. Please set it to True."
assert multitask, "independent GRPO training jobs for each predictor in the student program is not supported yet. Please set multitask=True."
# The backend will be called with a batch of (num_dspy_examples_per_grpo_step * num_rollouts_per_grpo_step * num_predictors) per training set if multitask is True
# If multitask is False, the backend will be called with a batch of (num_dspy_examples_per_grpo_step * num_rollouts_per_grpo_step) per training job
self.variably_invoked_predictor_grouping_mode = variably_invoked_predictor_grouping_mode
if variably_invoked_predictor_grouping_mode == "fill":
assert variably_invoked_predictor_fill_strategy is not None, "variably_invoked_predictor_fill_strategy must be set when variably_invoked_predictor_grouping_mode is 'fill'"
assert variably_invoked_predictor_fill_strategy in ["randint", "max"], "variably_invoked_predictor_fill_strategy must be either 'randint' or 'max'"
self.variably_invoked_predictor_fill_strategy = variably_invoked_predictor_fill_strategy
self.shuffled_trainset_ids = []
self.epoch = -1
self.id_freqs = Counter()
def validate_trace_data_and_log_issues(
self,
trace_data: list[list[list[dict[str, Any]]]],
subsample_training_dataset: list[Example],
num_teachers: int,
num_samples_per_input: int,
pred_signature_hash_to_ind: dict[int, int],
):
# At this point, trace_data: list[example_idx -> list[teacher_idx -> [num_samples_per_input * Dict(example, prediction, trace, example_ind, score)]]]
# Shape of trace is: [dspy_module_invocation_idx -> Tuple[Predictor, PredictorInputs, Prediction]]
assert len(trace_data) == len(subsample_training_dataset), f"Trace data length {len(trace_data)} does not match the number of examples {len(subsample_training_dataset)}"
assert len(trace_data[0]) == num_teachers, f"Trace data length {len(trace_data[0])} does not match the number of teachers {num_teachers}"
# TODO(GRPO Team): Ideally, once the dspy format issue is fixed, this change should be reverted back to being a normal assert.
if len(trace_data[0][0]) == 0:
logger.warning(f"Trace data for example {0} and teacher {0} is empty. This is likely due to all examples in the training set input, resulting in the model generating output not following the dspy response format.")
elif len(trace_data[0][0]) != num_samples_per_input:
logger.warning(f"Trace data length {len(trace_data[0][0])} does not match the expected number of samples per input {num_samples_per_input}")
assert "trace" in trace_data[0][0][0], "Trace data does not contain the 'trace' key"
assert len(trace_data[0][0][0]["trace"]) > 0, "Trace data is empty"
assert len(trace_data[0][0][0]["trace"][0]) == 3, f"Trace tuple length {len(trace_data[0][0][0]['trace'][0])} does not match the expected length 3"
for example_data in trace_data:
for teacher_data in example_data:
for sample in teacher_data:
for t in sample["trace"]:
assert hash(t[0].signature) in pred_signature_hash_to_ind
def report_validation_metrics(self, student, trainset, valset, logger, step_idx=-1):
if step_idx == -1 or step_idx == self.num_train_steps - 1 or (step_idx + 1) % self.num_steps_for_val == 0:
pass
else:
return
if valset is not None:
# Validation set provided by user
assert not self.use_train_as_val, "If valset is provided, use_train_as_val must be False."
assert isinstance(self.num_steps_for_val, int) and self.num_steps_for_val > 0, "num_steps_for_val must be a positive integer."
if self.report_train_scores:
if step_idx == -1:
logger.info("Using user provided validation set and reporting train scores for every validation step in addition.")
valset_evaluator = Evaluate(
devset=valset + trainset,
num_threads=self.num_threads,
display_progress=True,
provide_traceback=False, # TODO(check with team)
max_errors=len(valset)*10, # TODO(check with team)
failure_score=self.failure_score
)
if step_idx == -1:
logger.info("Evaluating the student program on the train+validation set before training loop...")
else:
logger.info(f"Evaluating the student program on the validation set after training step {step_idx + 1}/{self.num_train_steps}")
valset_evaluation = valset_evaluator(student, metric=self.metric)
trainset_scores = [r[-1] for r in valset_evaluation.results[len(valset):]]
valset_scores = [r[-1] for r in valset_evaluation.results[:len(valset)]]
trainset_agg = sum(trainset_scores) / len(trainset_scores)
valset_agg = sum(valset_scores) / len(valset_scores)
if step_idx == -1:
logger.info(f"Student program training set score before training loop: {trainset_agg}")
logger.info(f"Student program validation set score before training loop: {valset_agg}")
else:
logger.info(f"Student program training set score after training step {step_idx + 1}/{self.num_train_steps}: {trainset_agg}")
logger.info(f"Student program validation set score after training step {step_idx + 1}/{self.num_train_steps}: {valset_agg}")
else:
if step_idx == -1:
logger.info("Using user provided validation set and not reporting train scores.")
valset_evaluator = Evaluate(
devset=valset,
num_threads=self.num_threads,
display_progress=True,
provide_traceback=False, # TODO(check with team)
max_errors=len(valset)*10, # TODO(check with team)
failure_score=self.failure_score
)
if step_idx == -1:
logger.info("Evaluating the student program on the validation set before training loop...")
else:
logger.info(f"Evaluating the student program on the validation set after training step {step_idx + 1}/{self.num_train_steps}")
valset_evaluation = valset_evaluator(student, metric=self.metric)
if step_idx == -1:
logger.info(f"Student program validation set score before training loop: {valset_evaluation.score}")
else:
logger.info(f"Student program validation set score after training step {step_idx + 1}/{self.num_train_steps}: {valset_evaluation.score}")
else:
# No validation set provided by user
if self.report_train_scores:
assert self.use_train_as_val, "If report_train_scores is True, use_train_as_val must be True when valset is not provided explicitly."
assert isinstance(self.num_steps_for_val, int) and self.num_steps_for_val > 0, "num_steps_for_val must be a positive integer."
if step_idx == -1:
logger.info("Using trainset as validation set.")
valset_evaluator = Evaluate(
devset=trainset,
num_threads=self.num_threads,
display_progress=True,
provide_traceback=False, # TODO(check with team)
max_errors=len(trainset)*10, # TODO(check with team)
failure_score=self.failure_score
)
if step_idx == -1:
logger.info("Evaluating the student program on the validation set before training loop...")
else:
logger.info(f"Evaluating the student program on the validation set after training step {step_idx + 1}/{self.num_train_steps}")
valset_evaluation = valset_evaluator(student, metric=self.metric)
if step_idx == -1:
logger.info(f"Student program training set score before training loop: {valset_evaluation.score}")
else:
logger.info(f"Student program training set score after training step {step_idx + 1}/{self.num_train_steps}: {valset_evaluation.score}")
else:
# No valset provided, and not using train as val
assert not self.use_train_as_val, "If report_train_scores is False, use_train_as_val must be False."
if step_idx == -1:
logger.info("Not using any validation set and not reporting train scores.")
def update_shuffled_trainset(self, original_trainset):
self.shuffled_trainset_ids = list(range(len(original_trainset)))
self.rng.shuffle(self.shuffled_trainset_ids)
for id in self.shuffled_trainset_ids:
self.id_freqs[id] += 1
num_to_pad = self.num_dspy_examples_per_grpo_step - (len(original_trainset) % self.num_dspy_examples_per_grpo_step)
if num_to_pad > 0:
# Select ids based on least frequent ids
for _ in range(num_to_pad):
selected_id = self.id_freqs.most_common()[::-1][0][0]
self.shuffled_trainset_ids.append(selected_id)
self.id_freqs[selected_id] += 1
def select_training_sample_and_update_shuffled_trainset(
self,
original_trainset: list[Example],
train_step_idx: int,
) -> list[Example]:
base_idx = train_step_idx * self.num_dspy_examples_per_grpo_step
if self.epoch == -1:
curr_epoch = 0
else:
curr_epoch = base_idx // len(self.shuffled_trainset_ids)
if curr_epoch > self.epoch:
logger.info(f"Updating shuffled trainset for epoch {curr_epoch}...")
self.epoch = curr_epoch
self.update_shuffled_trainset(original_trainset)
assert len(self.shuffled_trainset_ids) >= self.num_dspy_examples_per_grpo_step, f"Shuffled trainset length {len(self.shuffled_trainset_ids)} is less than num_dspy_examples_per_grpo_step {self.num_dspy_examples_per_grpo_step}"
assert len(self.shuffled_trainset_ids) % self.num_dspy_examples_per_grpo_step == 0, f"Shuffled trainset length {len(self.shuffled_trainset_ids)} is not divisible by num_dspy_examples_per_grpo_step {self.num_dspy_examples_per_grpo_step}"
base_idx = base_idx % len(self.shuffled_trainset_ids)
end_idx = base_idx + self.num_dspy_examples_per_grpo_step
assert end_idx <= len(self.shuffled_trainset_ids), f"End index {end_idx} is out of bounds for shuffled trainset length {len(self.shuffled_trainset_ids)}"
selected_ids = self.shuffled_trainset_ids[base_idx:end_idx]
selected_trainset = [original_trainset[i] for i in selected_ids]
return selected_trainset
def compile(
self,
student: Module,
trainset: list[Example],
teacher: Module | list[Module] | None = None,
valset: list[Example] | None = None,
**kwargs,
) -> Module:
logger.info("Starting the GRPO compilation process... The LM(s) for the student program will be updated in place at the end of the training.")
logger.info("Validating the inputs...")
assert len(trainset) > 0, "Training set is empty. Please provide a non-empty training set."
if len(trainset) < self.num_dspy_examples_per_grpo_step:
logger.warning(
f"Number of training examples {len(trainset)} is less than the number of examples per GRPO step {self.num_dspy_examples_per_grpo_step}. "
"Repeating the training set to fill the GRPO step. This could lead to overfitting and training instability."
)
multiplier = (self.num_dspy_examples_per_grpo_step + len(trainset) - 1) // len(trainset)
if multiplier > 1:
logger.warning(
f"Repeating the training set {multiplier} times to fill the GRPO step. This could lead to overfitting and training instability."
)
trainset = trainset * multiplier
# TODO(GRPO Team): Following checks are for unimplemented features.
# Consider if we want to eventually implement them or remove. We don't
# yet support:
# * multitask == False
# * student program with multiple predictor LMs
# The main reason for these is that we update the LMs in place. If these
# LMs are shared between the different predictors of the student
# program and we have multitask == False, we need to decide which steps
# will use new LM copies and we need to ensure our decision is
# consistent with any teacher LMs that share the same LMs.
# TODO(GRPO Team): We want to make it possible to continue GRPO runs in
# the future by saving the state of the GRPO run in the event of a
# process failure.
if not self.multitask:
raise ValueError(
"Independent GRPO training jobs for each predictor in the student program "
"are not supported yet. Please set multitask=True."
)
student_lms = {id(pred.lm) for pred in student.predictors()}
assert len(student_lms) == 1, (
f"Student program has multiple LMs: {student_lms}. "
"GRPO only supports student programs with a single LM."
"You can set the LM for a program with `program.set_lm(...)`"
)
# Our regular input validation starts here
if self.use_train_as_val:
assert valset is None, "If use_train_as_val is True, valset must be None."
logger.info("Preparing the student program...")
all_predictors_have_lms(student)
pred_signature_hash_to_ind = {hash(pred.signature): ind for ind, pred in enumerate(student.predictors())}
num_student_predictors = len(student.predictors())
logging.info("Preparing the teacher program(s)... We will ensure that the provided programs have the same program structure as the student program.")
if (isinstance(teacher, list) and len(teacher) == 0) or teacher is None:
teacher = student
teachers = teacher if isinstance(teacher, list) else [teacher]
for t in teachers:
assert_structural_equivalency(student, t)
all_predictors_have_lms(t)
# Ensure that the teachers list contain the student program
assert student in teachers, f"Student program {student} is not in the list of teachers {teachers}. Please provide the student program as one of the teachers. Alternatively, you can leave the teacher argument as None, and the student program will be used as the teacher program."
assert self.num_rollouts_per_grpo_step % len(teachers) == 0, (
f"The GRPO group size (num_rollouts_per_grpo_step) {self.num_rollouts_per_grpo_step} is not divisible by the number of teachers {len(teachers)}. "
"This is required to ensure that each teacher gets the same number of examples."
"Please provide a number of examples that is divisible by the number of teachers."
)
num_samples_per_input = self.num_rollouts_per_grpo_step // len(teachers)
# We will disable the LM cache for all programs (student and teachers)
# These will be reverted to their original state at the end of the
# training
lm_cache_dict = {}
disable_lm_cache(program=student, lm_cache_dict=lm_cache_dict)
for t in teachers:
disable_lm_cache(program=t, lm_cache_dict=lm_cache_dict)
# Update train_kwargs
for pred in student.predictors():
train_kwargs = self.train_kwargs[pred.lm]
train_kwargs = {} if train_kwargs is None else train_kwargs
train_kwargs["num_generations"] = self.num_rollouts_per_grpo_step
self.train_kwargs[pred.lm] = train_kwargs
# We need to have a separate job for each unique LM x the data
# collection strategy. This properly handles all combinations of
# multitask and predictor LMs
logger.info("Preparing the GRPO training job(s)...")
grpo_training_jobs = {}
for pred_ind, pred in enumerate(student.predictors()):
data_key = None if self.multitask else pred_ind
job_key = (pred.lm, data_key)
if job_key not in grpo_training_jobs:
train_kwargs = self.train_kwargs[pred.lm]
job = pred.lm.reinforce(train_kwargs=train_kwargs, gpu_config=self.gpu_config)
grpo_training_jobs[job_key] = job
self.report_validation_metrics(
student=student,
trainset=trainset,
valset=valset,
logger=logger,
step_idx=-1,
)
logger.info("Starting the GRPO training loop...")
for train_step_idx in range(self.num_train_steps):
logger.info(f"GRPO training step {train_step_idx + 1}/{self.num_train_steps}...")
subsample_training_dataset = self.select_training_sample_and_update_shuffled_trainset(
original_trainset=trainset,
train_step_idx=train_step_idx,
)
logger.info("Bootstrapping data...")
trace_data = [[[] for _ in range(len(teachers))] for _ in range(len(subsample_training_dataset))]
for tind, teacher in enumerate(teachers):
subsample_training_dataset_repeated = [example for _ in range(num_samples_per_input) for example in subsample_training_dataset]
round_data = bootstrap_trace_data(
program=teacher,
dataset=subsample_training_dataset_repeated,
metric=self.metric,
num_threads=self.num_threads,
raise_on_error=False, # TODO(GRPO Team): This should be True, once the dspy format issue is fixed
capture_failed_parses=True,
failure_score=self.failure_score,
format_failure_score=self.format_failure_score,
)
for data_dict in round_data:
example_ind_in_subsample = data_dict["example_ind"] % len(subsample_training_dataset)
data_dict["example_ind"] = example_ind_in_subsample
trace_data[example_ind_in_subsample][tind].append(data_dict)
# The trace_data for examples with FailedPrediction cases will have the signature at index 0, instead of the predictor
# We need to replace the signature with the predictor
# At this point, trace_data: list[example_idx -> list[teacher_idx -> [num_samples_per_input * Dict(example, prediction, trace, example_ind, score)]]]
# Shape of trace is: [dspy_module_invocation_idx -> Tuple[Predictor, PredictorInputs, Prediction]]
self.validate_trace_data_and_log_issues(
trace_data=trace_data,
subsample_training_dataset=subsample_training_dataset,
num_teachers=len(teachers),
num_samples_per_input=num_samples_per_input,
pred_signature_hash_to_ind=pred_signature_hash_to_ind,
)
logger.info("Preparing the training data batch from bootstrapped examples for GRPO...")
# Now, we need to prepare batches of data to be sent for training
# Shape of train_batch_per_predictor: list[num_student_predictors -> list[ ]]
train_batch_per_predictor: list[list[GRPOGroup]] = [[] for _ in range(num_student_predictors)]
for pred_id in range(num_student_predictors):
for example_ind, example_data in enumerate(trace_data):
# Each example_data is a list of teacher_idx -> [num_samples_per_input * Dict(example, prediction, trace, example_ind, score)]
# We need to flatten this list and create a batch for each predictor
# TODO(Lakshya, Omar, Noah): Discuss what to do with the same module being invoked multiple times within a single dspy.Example
predictor_example_invocations: list[list[tuple]] = []
for teacher_data in example_data:
for sample in teacher_data:
# Each sample is a Dict(example, prediction, trace, example_ind, score)
# sample['prediction'] is module_level prediction
assert sample["example_ind"] == example_ind, f"Example index {sample['example_ind']} does not match the expected index {example_ind}"
trace_instances_for_current_pred = [(*t, sample["score"]) for t in sample["trace"] if hash(t[0].signature) == hash(student.predictors()[pred_id].signature)]
predictor_example_invocations.append(trace_instances_for_current_pred)
if len(predictor_example_invocations) == 0:
logger.warning(f"Skipping example {example_ind} for predictor {pred_id} as it has no invocations. This is likely due to all examples in the training set input, resulting in the model generating output not following the dspy response format.")
continue
elif len(predictor_example_invocations) != self.num_rollouts_per_grpo_step:
logger.warning(f"Number of predictor example invocations {len(predictor_example_invocations)} does not match the expected batch size {self.num_rollouts_per_grpo_step}. This is likely due to all examples in the training set input, resulting in the model generating output not following the dspy response format.")
min_len = min([len(predictor_example_invocations[i]) for i in range(len(predictor_example_invocations))])
max_len = max([len(predictor_example_invocations[i]) for i in range(len(predictor_example_invocations))])
if min_len == 0:
logger.warning(f"Skipping example {example_ind} for predictor {pred_id} as it has no invocations.")
continue
if self.variably_invoked_predictor_grouping_mode == "truncate":
predictor_example_invocations = [invocation[:min_len] for invocation in predictor_example_invocations]
elif self.variably_invoked_predictor_grouping_mode == "fill":
if self.variably_invoked_predictor_fill_strategy == "randint":
selector = lambda l: self.rng.choice(l) # noqa: E731, E741
else:
selector = lambda l: l[-1] # noqa: E731, E741
predictor_example_invocations = [
invocation + [selector(invocation) for _ in range(max_len - len(invocation))]
for invocation in predictor_example_invocations
]
else:
assert self.variably_invoked_predictor_grouping_mode == "ragged", f"Unknown variably invoked predictor grouping mode {self.variably_invoked_predictor_grouping_mode}"
max_len = max([len(predictor_example_invocations[i]) for i in range(len(predictor_example_invocations))])
example_training_data: list[GRPOGroup] = [[] for _ in range(max_len)]
for group_idx in range(max_len):
for rollout_idx in range(len(predictor_example_invocations)):
trace_instance = predictor_example_invocations[rollout_idx][group_idx]
score = trace_instance[3]
# for module_invocation_idx, trace_instance in enumerate(trace_instances_for_current_pred):
# Each trace is a tuple of (Predictor, PredictorInputs, Prediction)
trace_pred_id = pred_signature_hash_to_ind.get(hash(trace_instance[0].signature))
assert trace_pred_id == pred_id
predictor = trace_instance[0]
pred_lm = predictor.lm
adapter = self.adapter[pred_lm] or settings.adapter or ChatAdapter()
assert isinstance(adapter, ChatAdapter), f"Adapter {adapter} is not a ChatAdapter. GRPO training is not supported for this adapter."
# TODO(Lakshya): Currently we exclude demos from the training data
# TODO(GRPO Team): Use build_call_data_from_trace (from bootstrap_finetune) instead of
# dealing with the message formatting ourselves.
inp_messages = adapter.format(
signature=trace_instance[0].signature,
inputs=trace_instance[1],
demos=[] # TODO: Add support for demos
)
if isinstance(trace_instance[2], FailedPrediction):
score = trace_instance[2].format_reward or self.format_failure_score
example_training_data[group_idx].append({
"messages": inp_messages,
"completion": {
"role": "assistant",
"content": trace_instance[2].completion_text,
},
"reward": float(score),
})
logger.warning(f"Adding a format failure example to the training data for predictor {pred_id} and example {example_ind}.")
else:
all_messages = adapter.format_finetune_data(
signature=trace_instance[0].signature,
inputs=trace_instance[1],
outputs=trace_instance[2],
demos=[] # TODO: Add support for demos
)["messages"]
assert all_messages[:-1] == inp_messages, f"Input messages {inp_messages} do not match the expected messages {all_messages[:-1]}"
example_training_data[group_idx].append({
"messages": inp_messages,
"completion": {
"role": all_messages[-1]["role"],
"content": all_messages[-1]["content"],
},
"reward": float(score),
})
train_batch_per_predictor[pred_id].extend(example_training_data)
if not any(train_batch_per_predictor):
logger.warning("No training data found for this training step. This means that the model did not generate valid formatted responses for any of the examples in the training set. This is a critical error. Please check the model and the training set.")
continue
for predictor_train_batch in train_batch_per_predictor:
for grpo_train_group in predictor_train_batch:
if len(grpo_train_group) != self.num_rollouts_per_grpo_step:
logger.warning(f"Number of completions {len(grpo_train_group)} does not match the expected number num_rollouts_per_grpo_step={self.num_rollouts_per_grpo_step}")
assert len(grpo_train_group) <= self.num_rollouts_per_grpo_step, f"Number of completions {len(grpo_train_group)} is greater than the expected number num_rollouts_per_grpo_step={self.num_rollouts_per_grpo_step}"
if len(set(map(repr, grpo_train_group))) < 2:
# TODO(GRPO Team): How can we avoid this warning?
logger.warning(f"GRPOGroup has no diversity. This could be due to low temperature, or low number of rollouts, or the cache could be enabled inadvertently. The GRPOGroup is {grpo_train_group}.")
# We now run the GRPO step. Notes:
# * The job here has a reference to a particular M that's attached
# to the student program. We update the .model field of this LM
# inside the job, which also updates the LM in the student program
# since these point to the same reference (along with any teacher
# program that shares the same LM).
# * TODO(GRPO Team): This is inconsistent with how
# BootstrapFinetune works, which creates new LM instances post
# training. We should decide whether the LMs should be updated in
# place or new LMs should be created, and standardize our approach
# for both. If we decide to create new LMs, we should find a way
# to update self.adapter and self.train_kwargs accordingly, in
# addition to updating any teacher programs that share the same
# LM.
logger.info("Invoking GRPO training step...")
for (_, data_key), job in grpo_training_jobs.items():
train_data: list[GRPOGroup] = sum(train_batch_per_predictor, []) if data_key is None else train_batch_per_predictor[data_key] #noqa: RUF017
for group in train_data:
if len(group) != self.num_rollouts_per_grpo_step:
# TODO(GRPO Team): This is very undesirable. This occurs only because in some of the generations, the model does not follow the correct dspy format.
# The ideal solution is to identify the full response string in that predictor's group, and then assign a high-negative (user-configurable) reward to that group.
# Pad the group to the expected number of generations by repeating the whole group, might require multiple iterations
while len(group) < self.num_rollouts_per_grpo_step:
group.extend(group[:min(self.num_rollouts_per_grpo_step - len(group), len(group))])
assert len(group) == self.num_rollouts_per_grpo_step, f"Number of completions {len(group)} does not match the expected number self.num_rollouts_per_grpo_step={self.num_rollouts_per_grpo_step}"
job.step(train_data=train_data, train_data_format=TrainDataFormat.GRPO_CHAT)
logger.info(f"GRPO training step {train_step_idx + 1}/{self.num_train_steps} completed.")
self.report_validation_metrics(
student=student,
trainset=trainset,
valset=valset,
logger=logger,
step_idx=train_step_idx,
)
logger.info("Done with the iterations! Retrieving the final model(s)...")
for _, job in grpo_training_jobs.items():
job.terminate()
# Revert cache states to their initial values
recover_lm_cache(program=student, lm_cache_dict=lm_cache_dict)
for t in teachers:
recover_lm_cache(program=t, lm_cache_dict=lm_cache_dict)
logger.info("GRPO compiler has finished compiling the student program")
student._compiled = True
return student
def disable_lm_cache(program: Module, lm_cache_dict: dict):
"""Disable the LM cache for all predictors in the program."""
for pred in program.predictors():
if not pred.lm:
raise ValueError(f"Cannot disable cache: predictor {pred} does not have an LM set.")
if pred.lm not in lm_cache_dict: # Check to avoid overwriting the cache
lm_cache_dict[pred.lm] = pred.lm.cache
pred.lm.cache = False
def recover_lm_cache(program: Module, lm_cache_dict: dict):
"""Recover the LM caches for all predictors in the program to their original state."""
for pred in program.predictors():
if pred.lm in lm_cache_dict:
pred.lm.cache = lm_cache_dict[pred.lm]
else:
# We do not expect this branch to execute at all since all the LMs
# are modified in place and no new LMs are created during training.
# However, we do not complain if this happens since this is a
# relatively minor feature. We default the LM cache to True.
pred.lm.cache = True
```
--------------------------------------------------------------------------------
/tests/adapters/test_json_adapter.py:
--------------------------------------------------------------------------------
```python
from unittest import mock
import pydantic
import pytest
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
from litellm.utils import ChatCompletionMessageToolCall, Choices, Function, Message, ModelResponse
from openai.types.responses import ResponseOutputMessage
import dspy
def test_json_adapter_passes_structured_output_when_supported_by_model():
class OutputField3(pydantic.BaseModel):
subfield1: int = pydantic.Field(description="Int subfield 1", ge=0, le=10)
subfield2: float = pydantic.Field(description="Float subfield 2")
class TestSignature(dspy.Signature):
input1: str = dspy.InputField()
output1: str = dspy.OutputField() # Description intentionally left blank
output2: bool = dspy.OutputField(desc="Boolean output field")
output3: OutputField3 = dspy.OutputField(desc="Nested output field")
output4_unannotated = dspy.OutputField(desc="Unannotated output field")
program = dspy.Predict(TestSignature)
# Configure DSPy to use an OpenAI LM that supports structured outputs
dspy.configure(lm=dspy.LM(model="openai/gpt-4o"), adapter=dspy.JSONAdapter())
with mock.patch("litellm.completion") as mock_completion:
program(input1="Test input")
def clean_schema_extra(field_name, field_info):
attrs = dict(field_info.__repr_args__())
if "json_schema_extra" in attrs:
attrs["json_schema_extra"] = {
k: v
for k, v in attrs["json_schema_extra"].items()
if k != "__dspy_field_type" and not (k == "desc" and v == f"${{{field_name}}}")
}
return attrs
mock_completion.assert_called_once()
_, call_kwargs = mock_completion.call_args
response_format = call_kwargs.get("response_format")
assert response_format is not None
assert issubclass(response_format, pydantic.BaseModel)
assert response_format.model_fields.keys() == {"output1", "output2", "output3", "output4_unannotated"}
def test_json_adapter_not_using_structured_outputs_when_not_supported_by_model():
class TestSignature(dspy.Signature):
input1: str = dspy.InputField()
output1: str = dspy.OutputField()
output2: bool = dspy.OutputField()
program = dspy.Predict(TestSignature)
# Configure DSPy to use a model from a fake provider that doesn't support structured outputs
dspy.configure(lm=dspy.LM(model="fakeprovider/fakemodel", cache=False), adapter=dspy.JSONAdapter())
with mock.patch("litellm.completion") as mock_completion:
mock_completion.return_value = ModelResponse(
choices=[Choices(message=Message(content=("{'output1': 'Test output', 'output2': True}")))],
model="openai/gpt-4o",
)
program(input1="Test input")
mock_completion.assert_called_once()
_, call_kwargs = mock_completion.call_args
assert "response_format" not in call_kwargs
def test_json_adapter_falls_back_when_structured_outputs_fails():
class TestSignature(dspy.Signature):
input1: str = dspy.InputField()
output1: str = dspy.OutputField(desc="String output field")
dspy.configure(lm=dspy.LM(model="openai/gpt-4o", cache=False), adapter=dspy.JSONAdapter())
program = dspy.Predict(TestSignature)
with mock.patch("litellm.completion") as mock_completion:
mock_completion.side_effect = [Exception("Bad structured outputs!"), mock_completion.return_value]
program(input1="Test input")
assert mock_completion.call_count == 2
_, first_call_kwargs = mock_completion.call_args_list[0]
assert issubclass(first_call_kwargs.get("response_format"), pydantic.BaseModel)
_, second_call_kwargs = mock_completion.call_args_list[1]
assert second_call_kwargs.get("response_format") == {"type": "json_object"}
def test_json_adapter_with_structured_outputs_does_not_mutate_original_signature():
class OutputField3(pydantic.BaseModel):
subfield1: int = pydantic.Field(description="Int subfield 1")
subfield2: float = pydantic.Field(description="Float subfield 2")
class TestSignature(dspy.Signature):
input1: str = dspy.InputField()
output1: str = dspy.OutputField() # Description intentionally left blank
output2: bool = dspy.OutputField(desc="Boolean output field")
output3: OutputField3 = dspy.OutputField(desc="Nested output field")
output4_unannotated = dspy.OutputField(desc="Unannotated output field")
dspy.configure(lm=dspy.LM(model="openai/gpt-4o"), adapter=dspy.JSONAdapter())
program = dspy.Predict(TestSignature)
with mock.patch("litellm.completion"):
program(input1="Test input")
assert program.signature.output_fields == TestSignature.output_fields
def test_json_adapter_sync_call():
signature = dspy.make_signature("question->answer")
adapter = dspy.JSONAdapter()
lm = dspy.utils.DummyLM([{"answer": "Paris"}], adapter=adapter)
with dspy.context(adapter=adapter):
result = adapter(lm, {}, signature, [], {"question": "What is the capital of France?"})
assert result == [{"answer": "Paris"}]
@pytest.mark.asyncio
async def test_json_adapter_async_call():
signature = dspy.make_signature("question->answer")
adapter = dspy.JSONAdapter()
lm = dspy.utils.DummyLM([{"answer": "Paris"}], adapter=adapter)
with dspy.context(adapter=adapter):
result = await adapter.acall(lm, {}, signature, [], {"question": "What is the capital of France?"})
assert result == [{"answer": "Paris"}]
def test_json_adapter_on_pydantic_model():
from litellm.utils import Choices, Message, ModelResponse
class User(pydantic.BaseModel):
id: int
name: str
email: str
class Answer(pydantic.BaseModel):
analysis: str
result: str
class TestSignature(dspy.Signature):
user: User = dspy.InputField(desc="The user who asks the question")
question: str = dspy.InputField(desc="Question the user asks")
answer: Answer = dspy.OutputField(desc="Answer to this question")
program = dspy.Predict(TestSignature)
dspy.configure(lm=dspy.LM(model="openai/gpt-4o", cache=False), adapter=dspy.JSONAdapter())
with mock.patch("litellm.completion") as mock_completion:
mock_completion.return_value = ModelResponse(
choices=[
Choices(
message=Message(
content="{'answer': {'analysis': 'Paris is the capital of France', 'result': 'Paris'}}"
)
)
],
model="openai/gpt-4o",
)
result = program(
user={"id": 5, "name": "name_test", "email": "email_test"}, question="What is the capital of France?"
)
# Check that litellm.completion was called exactly once
mock_completion.assert_called_once()
_, call_kwargs = mock_completion.call_args
# Assert that there are exactly 2 messages (system + user)
assert len(call_kwargs["messages"]) == 2
assert call_kwargs["messages"][0]["role"] == "system"
content = call_kwargs["messages"][0]["content"]
assert content is not None
# Assert that system prompt includes correct input field descriptions
expected_input_fields = (
"1. `user` (User): The user who asks the question\n2. `question` (str): Question the user asks\n"
)
assert expected_input_fields in content
# Assert that system prompt includes correct output field description
expected_output_fields = "1. `answer` (Answer): Answer to this question\n"
assert expected_output_fields in content
# Assert that system prompt includes input formatting structure
expected_input_structure = "[[ ## user ## ]]\n{user}\n\n[[ ## question ## ]]\n{question}\n\n"
assert expected_input_structure in content
# Assert that system prompt includes output formatting structure
expected_output_structure = (
"Outputs will be a JSON object with the following fields.\n\n{\n "
'"answer": "{answer} # note: the value you produce must adhere to the JSON schema: '
'{\\"type\\": \\"object\\", \\"properties\\": {\\"analysis\\": {\\"type\\": \\"string\\", \\"title\\": '
'\\"Analysis\\"}, \\"result\\": {\\"type\\": \\"string\\", \\"title\\": \\"Result\\"}}, \\"required\\": '
'[\\"analysis\\", \\"result\\"], \\"title\\": \\"Answer\\"}"\n}'
)
assert expected_output_structure in content
assert call_kwargs["messages"][1]["role"] == "user"
user_message_content = call_kwargs["messages"][1]["content"]
assert user_message_content is not None
# Assert that the user input data is formatted correctly
expected_input_data = (
'[[ ## user ## ]]\n{"id": 5, "name": "name_test", "email": "email_test"}\n\n[[ ## question ## ]]\n'
"What is the capital of France?\n\n"
)
assert expected_input_data in user_message_content
# Assert that the adapter output has expected fields and values
assert result.answer.analysis == "Paris is the capital of France"
assert result.answer.result == "Paris"
def test_json_adapter_parse_raise_error_on_mismatch_fields():
signature = dspy.make_signature("question->answer")
adapter = dspy.JSONAdapter()
with mock.patch("litellm.completion") as mock_completion:
mock_completion.return_value = ModelResponse(
choices=[
Choices(message=Message(content="{'answer1': 'Paris'}")),
],
model="openai/gpt-4o",
)
lm = dspy.LM(model="openai/gpt-4o-mini")
with pytest.raises(dspy.utils.exceptions.AdapterParseError) as e:
adapter(lm, {}, signature, [], {"question": "What is the capital of France?"})
assert e.value.adapter_name == "JSONAdapter"
assert e.value.signature == signature
assert e.value.lm_response == "{'answer1': 'Paris'}"
assert e.value.parsed_result == {}
assert str(e.value) == (
"Adapter JSONAdapter failed to parse the LM response. \n\n"
"LM Response: {'answer1': 'Paris'} \n\n"
"Expected to find output fields in the LM response: [answer] \n\n"
"Actual output fields parsed from the LM response: [] \n\n"
)
def test_json_adapter_formats_image():
# Test basic image formatting
image = dspy.Image(url="https://example.com/image.jpg")
class MySignature(dspy.Signature):
image: dspy.Image = dspy.InputField()
text: str = dspy.OutputField()
adapter = dspy.JSONAdapter()
messages = adapter.format(MySignature, [], {"image": image})
assert len(messages) == 2
user_message_content = messages[1]["content"]
assert user_message_content is not None
# The message should have 3 chunks of types: text, image_url, text
assert len(user_message_content) == 3
assert user_message_content[0]["type"] == "text"
assert user_message_content[2]["type"] == "text"
# Assert that the image is formatted correctly
expected_image_content = {"type": "image_url", "image_url": {"url": "https://example.com/image.jpg"}}
assert expected_image_content in user_message_content
def test_json_adapter_formats_image_with_few_shot_examples():
class MySignature(dspy.Signature):
image: dspy.Image = dspy.InputField()
text: str = dspy.OutputField()
adapter = dspy.JSONAdapter()
demos = [
dspy.Example(
image=dspy.Image(url="https://example.com/image1.jpg"),
text="This is a test image",
),
dspy.Example(
image=dspy.Image(url="https://example.com/image2.jpg"),
text="This is another test image",
),
]
messages = adapter.format(MySignature, demos, {"image": dspy.Image(url="https://example.com/image3.jpg")})
# 1 system message, 2 few shot examples (1 user and assistant message for each example), 1 user message
assert len(messages) == 6
assert {"type": "image_url", "image_url": {"url": "https://example.com/image1.jpg"}} in messages[1]["content"]
assert {"type": "image_url", "image_url": {"url": "https://example.com/image2.jpg"}} in messages[3]["content"]
assert {"type": "image_url", "image_url": {"url": "https://example.com/image3.jpg"}} in messages[5]["content"]
def test_json_adapter_formats_image_with_nested_images():
class ImageWrapper(pydantic.BaseModel):
images: list[dspy.Image]
tag: list[str]
class MySignature(dspy.Signature):
image: ImageWrapper = dspy.InputField()
text: str = dspy.OutputField()
image1 = dspy.Image(url="https://example.com/image1.jpg")
image2 = dspy.Image(url="https://example.com/image2.jpg")
image3 = dspy.Image(url="https://example.com/image3.jpg")
image_wrapper = ImageWrapper(images=[image1, image2, image3], tag=["test", "example"])
adapter = dspy.JSONAdapter()
messages = adapter.format(MySignature, [], {"image": image_wrapper})
expected_image1_content = {"type": "image_url", "image_url": {"url": "https://example.com/image1.jpg"}}
expected_image2_content = {"type": "image_url", "image_url": {"url": "https://example.com/image2.jpg"}}
expected_image3_content = {"type": "image_url", "image_url": {"url": "https://example.com/image3.jpg"}}
assert expected_image1_content in messages[1]["content"]
assert expected_image2_content in messages[1]["content"]
assert expected_image3_content in messages[1]["content"]
def test_json_adapter_formats_with_nested_documents():
class DocumentWrapper(pydantic.BaseModel):
documents: list[dspy.experimental.Document]
class MySignature(dspy.Signature):
document: DocumentWrapper = dspy.InputField()
text: str = dspy.OutputField()
doc1 = dspy.experimental.Document(data="Hello, world!")
doc2 = dspy.experimental.Document(data="Hello, world 2!")
document_wrapper = DocumentWrapper(documents=[doc1, doc2])
adapter = dspy.JSONAdapter()
messages = adapter.format(MySignature, [], {"document": document_wrapper})
expected_doc1_content = {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "Hello, world!"}, "citations": {"enabled": True}}
expected_doc2_content = {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "Hello, world 2!"}, "citations": {"enabled": True}}
assert expected_doc1_content in messages[1]["content"]
assert expected_doc2_content in messages[1]["content"]
def test_json_adapter_formats_image_with_few_shot_examples_with_nested_images():
class ImageWrapper(pydantic.BaseModel):
images: list[dspy.Image]
tag: list[str]
class MySignature(dspy.Signature):
image: ImageWrapper = dspy.InputField()
text: str = dspy.OutputField()
image1 = dspy.Image(url="https://example.com/image1.jpg")
image2 = dspy.Image(url="https://example.com/image2.jpg")
image3 = dspy.Image(url="https://example.com/image3.jpg")
image_wrapper = ImageWrapper(images=[image1, image2, image3], tag=["test", "example"])
demos = [
dspy.Example(
image=image_wrapper,
text="This is a test image",
),
]
image_wrapper_2 = ImageWrapper(images=[dspy.Image(url="https://example.com/image4.jpg")], tag=["test", "example"])
adapter = dspy.JSONAdapter()
messages = adapter.format(MySignature, demos, {"image": image_wrapper_2})
assert len(messages) == 4
# Image information in the few-shot example's user message
expected_image1_content = {"type": "image_url", "image_url": {"url": "https://example.com/image1.jpg"}}
expected_image2_content = {"type": "image_url", "image_url": {"url": "https://example.com/image2.jpg"}}
expected_image3_content = {"type": "image_url", "image_url": {"url": "https://example.com/image3.jpg"}}
assert expected_image1_content in messages[1]["content"]
assert expected_image2_content in messages[1]["content"]
assert expected_image3_content in messages[1]["content"]
# The query image is formatted in the last user message
assert {"type": "image_url", "image_url": {"url": "https://example.com/image4.jpg"}} in messages[-1]["content"]
def test_json_adapter_with_tool():
class MySignature(dspy.Signature):
"""Answer question with the help of the tools"""
question: str = dspy.InputField()
tools: list[dspy.Tool] = dspy.InputField()
answer: str = dspy.OutputField()
tool_calls: dspy.ToolCalls = dspy.OutputField()
def get_weather(city: str) -> str:
"""Get the weather for a city"""
return f"The weather in {city} is sunny"
def get_population(country: str, year: int) -> str:
"""Get the population for a country"""
return f"The population of {country} in {year} is 1000000"
tools = [dspy.Tool(get_weather), dspy.Tool(get_population)]
adapter = dspy.JSONAdapter()
messages = adapter.format(MySignature, [], {"question": "What is the weather in Tokyo?", "tools": tools})
assert len(messages) == 2
# The output field type description should be included in the system message even if the output field is nested
assert dspy.ToolCalls.description() in messages[0]["content"]
# The user message should include the question and the tools
assert "What is the weather in Tokyo?" in messages[1]["content"]
assert "get_weather" in messages[1]["content"]
assert "get_population" in messages[1]["content"]
# Tool arguments format should be included in the user message
assert "{'city': {'type': 'string'}}" in messages[1]["content"]
assert "{'country': {'type': 'string'}, 'year': {'type': 'integer'}}" in messages[1]["content"]
with mock.patch("litellm.completion") as mock_completion:
lm = dspy.LM(model="openai/gpt-4o-mini")
adapter(lm, {}, MySignature, [], {"question": "What is the weather in Tokyo?", "tools": tools})
mock_completion.assert_called_once()
_, call_kwargs = mock_completion.call_args
# Assert tool calls are included in the `tools` arg
assert len(call_kwargs["tools"]) > 0
assert call_kwargs["tools"][0] == {
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the weather for a city",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
},
},
"required": ["city"],
},
},
}
assert call_kwargs["tools"][1] == {
"type": "function",
"function": {
"name": "get_population",
"description": "Get the population for a country",
"parameters": {
"type": "object",
"properties": {
"country": {
"type": "string",
},
"year": {
"type": "integer",
},
},
"required": ["country", "year"],
},
},
}
def test_json_adapter_with_code():
# Test with code as input field
class CodeAnalysis(dspy.Signature):
"""Analyze the time complexity of the code"""
code: dspy.Code = dspy.InputField()
result: str = dspy.OutputField()
adapter = dspy.JSONAdapter()
messages = adapter.format(CodeAnalysis, [], {"code": "print('Hello, world!')"})
assert len(messages) == 2
# The output field type description should be included in the system message even if the output field is nested
assert dspy.Code.description() in messages[0]["content"]
# The user message should include the question and the tools
assert "print('Hello, world!')" in messages[1]["content"]
# Test with code as output field
class CodeGeneration(dspy.Signature):
"""Generate code to answer the question"""
question: str = dspy.InputField()
code: dspy.Code = dspy.OutputField()
adapter = dspy.JSONAdapter()
with mock.patch("litellm.completion") as mock_completion:
mock_completion.return_value = ModelResponse(
choices=[Choices(message=Message(content="{'code': 'print(\"Hello, world!\")'}"))],
model="openai/gpt-4o-mini",
)
result = adapter(
dspy.LM(model="openai/gpt-4o-mini", cache=False),
{},
CodeGeneration,
[],
{"question": "Write a python program to print 'Hello, world!'"},
)
assert result[0]["code"].code == 'print("Hello, world!")'
def test_json_adapter_formats_conversation_history():
class MySignature(dspy.Signature):
question: str = dspy.InputField()
history: dspy.History = dspy.InputField()
answer: str = dspy.OutputField()
history = dspy.History(
messages=[
{"question": "What is the capital of France?", "answer": "Paris"},
{"question": "What is the capital of Germany?", "answer": "Berlin"},
]
)
adapter = dspy.JSONAdapter()
messages = adapter.format(MySignature, [], {"question": "What is the capital of France?", "history": history})
assert len(messages) == 6
assert messages[1]["content"] == "[[ ## question ## ]]\nWhat is the capital of France?"
assert messages[2]["content"] == '{\n "answer": "Paris"\n}'
assert messages[3]["content"] == "[[ ## question ## ]]\nWhat is the capital of Germany?"
assert messages[4]["content"] == '{\n "answer": "Berlin"\n}'
@pytest.mark.asyncio
async def test_json_adapter_on_pydantic_model_async():
from litellm.utils import Choices, Message, ModelResponse
class User(pydantic.BaseModel):
id: int
name: str
email: str
class Answer(pydantic.BaseModel):
analysis: str
result: str
class TestSignature(dspy.Signature):
user: User = dspy.InputField(desc="The user who asks the question")
question: str = dspy.InputField(desc="Question the user asks")
answer: Answer = dspy.OutputField(desc="Answer to this question")
program = dspy.Predict(TestSignature)
with mock.patch("litellm.acompletion") as mock_completion:
mock_completion.return_value = ModelResponse(
choices=[
Choices(
message=Message(
content="{'answer': {'analysis': 'Paris is the capital of France', 'result': 'Paris'}}"
)
)
],
model="openai/gpt-4o",
)
with dspy.context(lm=dspy.LM(model="openai/gpt-4o", cache=False), adapter=dspy.JSONAdapter()):
result = await program.acall(
user={"id": 5, "name": "name_test", "email": "email_test"}, question="What is the capital of France?"
)
# Check that litellm.acompletion was called exactly once
mock_completion.assert_called_once()
_, call_kwargs = mock_completion.call_args
# Assert that there are exactly 2 messages (system + user)
assert len(call_kwargs["messages"]) == 2
assert call_kwargs["messages"][0]["role"] == "system"
content = call_kwargs["messages"][0]["content"]
assert content is not None
# Assert that system prompt includes correct input field descriptions
expected_input_fields = (
"1. `user` (User): The user who asks the question\n2. `question` (str): Question the user asks\n"
)
assert expected_input_fields in content
# Assert that system prompt includes correct output field description
expected_output_fields = "1. `answer` (Answer): Answer to this question\n"
assert expected_output_fields in content
# Assert that system prompt includes input formatting structure
expected_input_structure = "[[ ## user ## ]]\n{user}\n\n[[ ## question ## ]]\n{question}\n\n"
assert expected_input_structure in content
# Assert that system prompt includes output formatting structure
expected_output_structure = (
"Outputs will be a JSON object with the following fields.\n\n{\n "
'"answer": "{answer} # note: the value you produce must adhere to the JSON schema: '
'{\\"type\\": \\"object\\", \\"properties\\": {\\"analysis\\": {\\"type\\": \\"string\\", \\"title\\": '
'\\"Analysis\\"}, \\"result\\": {\\"type\\": \\"string\\", \\"title\\": \\"Result\\"}}, \\"required\\": '
'[\\"analysis\\", \\"result\\"], \\"title\\": \\"Answer\\"}"\n}'
)
assert expected_output_structure in content
assert call_kwargs["messages"][1]["role"] == "user"
user_message_content = call_kwargs["messages"][1]["content"]
assert user_message_content is not None
# Assert that the user input data is formatted correctly
expected_input_data = (
'[[ ## user ## ]]\n{"id": 5, "name": "name_test", "email": "email_test"}\n\n[[ ## question ## ]]\n'
"What is the capital of France?\n\n"
)
assert expected_input_data in user_message_content
# Assert that the adapter output has expected fields and values
assert result.answer.analysis == "Paris is the capital of France"
assert result.answer.result == "Paris"
def test_json_adapter_fallback_to_json_mode_on_structured_output_failure():
class TestSignature(dspy.Signature):
question: str = dspy.InputField()
answer: str = dspy.OutputField(desc="String output field")
dspy.configure(lm=dspy.LM(model="openai/gpt-4o-mini", cache=False), adapter=dspy.JSONAdapter())
program = dspy.Predict(TestSignature)
with mock.patch("litellm.completion") as mock_completion:
# First call raises error to simulate structured output failure, second call returns a valid response
mock_completion.side_effect = [
RuntimeError("Structured output failed!"),
ModelResponse(choices=[Choices(message=Message(content="{'answer': 'Test output'}"))]),
]
result = program(question="Dummy question!")
# The parse should succeed on the second call
assert mock_completion.call_count == 2
assert result.answer == "Test output"
# The first call should have tried structured output
_, first_call_kwargs = mock_completion.call_args_list[0]
assert issubclass(first_call_kwargs.get("response_format"), pydantic.BaseModel)
# The second call should have used JSON mode
_, second_call_kwargs = mock_completion.call_args_list[1]
assert second_call_kwargs.get("response_format") == {"type": "json_object"}
def test_json_adapter_json_mode_no_structured_outputs():
class TestSignature(dspy.Signature):
question: str = dspy.InputField()
answer: str = dspy.OutputField(desc="String output field")
dspy.configure(lm=dspy.LM(model="openai/gpt-4o", cache=False), adapter=dspy.JSONAdapter())
program = dspy.Predict(TestSignature)
with mock.patch("litellm.completion") as mock_completion, \
mock.patch("litellm.get_supported_openai_params") as mock_get_supported_openai_params, \
mock.patch("litellm.supports_response_schema") as mock_supports_response_schema:
# Call a model that allows json but not structured outputs
mock_completion.return_value = ModelResponse(choices=[Choices(message=Message(content="{'answer': 'Test output'}"))])
mock_get_supported_openai_params.return_value = ["response_format"]
mock_supports_response_schema.return_value = False
result = program(question="Dummy question!")
assert mock_completion.call_count == 1
assert result.answer == "Test output"
_, call_kwargs = mock_completion.call_args_list[0]
assert call_kwargs.get("response_format") == {"type": "json_object"}
@pytest.mark.asyncio
async def test_json_adapter_json_mode_no_structured_outputs_async():
class TestSignature(dspy.Signature):
question: str = dspy.InputField()
answer: str = dspy.OutputField(desc="String output field")
program = dspy.Predict(TestSignature)
with mock.patch("litellm.acompletion") as mock_acompletion, \
mock.patch("litellm.get_supported_openai_params") as mock_get_supported_openai_params, \
mock.patch("litellm.supports_response_schema") as mock_supports_response_schema:
# Call a model that allows json but not structured outputs
mock_acompletion.return_value = ModelResponse(choices=[Choices(message=Message(content="{'answer': 'Test output'}"))])
mock_get_supported_openai_params.return_value = ["response_format"]
mock_supports_response_schema.return_value = False
with dspy.context(lm=dspy.LM(model="openai/gpt-4o", cache=False), adapter=dspy.JSONAdapter()):
result = await program.acall(question="Dummy question!")
assert mock_acompletion.call_count == 1
assert result.answer == "Test output"
_, call_kwargs = mock_acompletion.call_args_list[0]
assert call_kwargs.get("response_format") == {"type": "json_object"}
@pytest.mark.asyncio
async def test_json_adapter_fallback_to_json_mode_on_structured_output_failure_async():
class TestSignature(dspy.Signature):
question: str = dspy.InputField()
answer: str = dspy.OutputField(desc="String output field")
program = dspy.Predict(TestSignature)
with mock.patch("litellm.acompletion") as mock_acompletion:
# First call raises error to simulate structured output failure, second call returns a valid response
mock_acompletion.side_effect = [
RuntimeError("Structured output failed!"),
ModelResponse(choices=[Choices(message=Message(content="{'answer': 'Test output'}"))]),
]
with dspy.context(lm=dspy.LM(model="openai/gpt-4o-mini", cache=False), adapter=dspy.JSONAdapter()):
result = await program.acall(question="Dummy question!")
# The parse should succeed on the second call
assert mock_acompletion.call_count == 2
assert result.answer == "Test output"
# The first call should have tried structured output
_, first_call_kwargs = mock_acompletion.call_args_list[0]
assert issubclass(first_call_kwargs.get("response_format"), pydantic.BaseModel)
# The second call should have used JSON mode
_, second_call_kwargs = mock_acompletion.call_args_list[1]
assert second_call_kwargs.get("response_format") == {"type": "json_object"}
def test_error_message_on_json_adapter_failure():
class TestSignature(dspy.Signature):
question: str = dspy.InputField()
answer: str = dspy.OutputField(desc="String output field")
program = dspy.Predict(TestSignature)
dspy.configure(lm=dspy.LM(model="openai/gpt-4o-mini", cache=False), adapter=dspy.JSONAdapter())
with mock.patch("litellm.completion") as mock_completion:
mock_completion.side_effect = RuntimeError("RuntimeError!")
with pytest.raises(RuntimeError) as error:
program(question="Dummy question!")
assert "RuntimeError!" in str(error.value)
mock_completion.side_effect = ValueError("ValueError!")
with pytest.raises(ValueError) as error:
program(question="Dummy question!")
assert "ValueError!" in str(error.value)
@pytest.mark.asyncio
async def test_error_message_on_json_adapter_failure_async():
class TestSignature(dspy.Signature):
question: str = dspy.InputField()
answer: str = dspy.OutputField(desc="String output field")
program = dspy.Predict(TestSignature)
with mock.patch("litellm.acompletion") as mock_acompletion:
with dspy.context(lm=dspy.LM(model="openai/gpt-4o-mini", cache=False), adapter=dspy.JSONAdapter()):
mock_acompletion.side_effect = RuntimeError("RuntimeError!")
with pytest.raises(RuntimeError) as error:
await program.acall(question="Dummy question!")
assert "RuntimeError!" in str(error.value)
mock_acompletion.side_effect = ValueError("ValueError!")
with pytest.raises(ValueError) as error:
await program.acall(question="Dummy question!")
assert "ValueError!" in str(error.value)
def test_json_adapter_toolcalls_native_function_calling():
class MySignature(dspy.Signature):
question: str = dspy.InputField()
tools: list[dspy.Tool] = dspy.InputField()
answer: str = dspy.OutputField()
tool_calls: dspy.ToolCalls = dspy.OutputField()
def get_weather(city: str) -> str:
return f"The weather in {city} is sunny"
tools = [dspy.Tool(get_weather)]
adapter = dspy.JSONAdapter(use_native_function_calling=True)
# Case 1: Tool calls are present in the response, while content is None.
with mock.patch("litellm.completion") as mock_completion:
mock_completion.return_value = ModelResponse(
choices=[
Choices(
finish_reason="tool_calls",
index=0,
message=Message(
content=None,
role="assistant",
tool_calls=[
ChatCompletionMessageToolCall(
function=Function(arguments='{"city":"Paris"}', name="get_weather"),
id="call_pQm8ajtSMxgA0nrzK2ivFmxG",
type="function",
)
],
),
),
],
model="openai/gpt-4o-mini",
)
result = adapter(
dspy.LM(model="openai/gpt-4o-mini", cache=False),
{},
MySignature,
[],
{"question": "What is the weather in Paris?", "tools": tools},
)
assert result[0]["tool_calls"] == dspy.ToolCalls(
tool_calls=[dspy.ToolCalls.ToolCall(name="get_weather", args={"city": "Paris"})]
)
# `answer` is not present, so we set it to None
assert result[0]["answer"] is None
# Case 2: Tool calls are not present in the response, while content is present.
with mock.patch("litellm.completion") as mock_completion:
mock_completion.return_value = ModelResponse(
choices=[Choices(message=Message(content="{'answer': 'Paris'}"))],
model="openai/gpt-4o-mini",
)
result = adapter(
dspy.LM(model="openai/gpt-4o-mini", cache=False),
{},
MySignature,
[],
{"question": "What is the weather in Paris?", "tools": tools},
)
assert result[0]["answer"] == "Paris"
assert result[0]["tool_calls"] is None
def test_json_adapter_toolcalls_no_native_function_calling():
class MySignature(dspy.Signature):
question: str = dspy.InputField()
tools: list[dspy.Tool] = dspy.InputField()
answer: str = dspy.OutputField()
tool_calls: dspy.ToolCalls = dspy.OutputField()
def get_weather(city: str) -> str:
return f"The weather in {city} is sunny"
tools = [dspy.Tool(get_weather)]
# Patch _get_structured_outputs_response_format to track calls
with mock.patch("dspy.adapters.json_adapter._get_structured_outputs_response_format") as mock_structured:
# Patch litellm.completion to return a dummy response
with mock.patch("litellm.completion") as mock_completion:
mock_completion.return_value = ModelResponse(
choices=[Choices(message=Message(content="{'answer': 'sunny', 'tool_calls': {'tool_calls': []}}"))],
model="openai/gpt-4o-mini",
)
adapter = dspy.JSONAdapter(use_native_function_calling=False)
lm = dspy.LM(model="openai/gpt-4o-mini", cache=False)
adapter(lm, {}, MySignature, [], {"question": "What is the weather in Tokyo?", "tools": tools})
# _get_structured_outputs_response_format is not called because without using native function calling,
# JSONAdapter falls back to json mode for stable quality.
mock_structured.assert_not_called()
mock_completion.assert_called_once()
_, call_kwargs = mock_completion.call_args
assert call_kwargs["response_format"] == {"type": "json_object"}
def test_json_adapter_with_responses_api():
class TestSignature(dspy.Signature):
question: str = dspy.InputField()
answer: str = dspy.OutputField()
api_response = ResponsesAPIResponse(
id="resp_1",
created_at=0.0,
error=None,
incomplete_details=None,
instructions=None,
model="openai/gpt-4o",
object="response",
output=[
ResponseOutputMessage(
**{
"id": "msg_1",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [
{"type": "output_text", "text": '{"answer": "Washington, D.C."}', "annotations": []}
],
},
),
],
metadata={},
parallel_tool_calls=False,
temperature=1.0,
tool_choice="auto",
tools=[],
top_p=1.0,
max_output_tokens=None,
previous_response_id=None,
reasoning=None,
status="completed",
text=None,
truncation="disabled",
usage=ResponseAPIUsage(input_tokens=10, output_tokens=5, total_tokens=15),
user=None,
)
lm = dspy.LM(model="openai/gpt-4o", model_type="responses", cache=False)
dspy.configure(lm=lm, adapter=dspy.JSONAdapter())
program = dspy.Predict(TestSignature)
with mock.patch("litellm.responses", autospec=True, return_value=api_response) as mock_responses:
result = program(question="What is the capital of the USA?")
assert result.answer == "Washington, D.C."
mock_responses.assert_called_once()
# Verify that response_format was converted to text.format
call_kwargs = mock_responses.call_args.kwargs
assert "response_format" not in call_kwargs
assert "text" in call_kwargs
assert isinstance(call_kwargs["text"]["format"], type)
assert issubclass(call_kwargs["text"]["format"], pydantic.BaseModel)
```
--------------------------------------------------------------------------------
/tests/streaming/test_streaming.py:
--------------------------------------------------------------------------------
```python
import asyncio
import time
from dataclasses import dataclass
from unittest import mock
from unittest.mock import AsyncMock
import pytest
from asyncer import syncify
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
import dspy
from dspy.adapters.types import Type
from dspy.experimental import Citations, Document
from dspy.streaming import StatusMessage, StatusMessageProvider, streaming_response
@pytest.mark.anyio
async def test_streamify_yields_expected_response_chunks(litellm_test_server):
api_base, _ = litellm_test_server
lm = dspy.LM(
model="openai/dspy-test-model",
api_base=api_base,
api_key="fakekey",
cache=True,
)
with dspy.context(lm=lm, adapter=dspy.JSONAdapter()):
class TestSignature(dspy.Signature):
input_text: str = dspy.InputField()
output_text: str = dspy.OutputField()
program = dspy.streamify(dspy.Predict(TestSignature))
output_stream1 = program(input_text="Test")
output_chunks1 = [chunk async for chunk in output_stream1]
last_chunk1 = output_chunks1[-1]
assert isinstance(last_chunk1, dspy.Prediction)
assert last_chunk1.output_text == "Hello!"
output_stream2 = program(input_text="Test")
output_chunks2 = [chunk async for chunk in output_stream2]
# Since the input is cached, only one chunk should be
# yielded containing the prediction
assert len(output_chunks2) == 1
last_chunk2 = output_chunks2[-1]
assert isinstance(last_chunk2, dspy.Prediction)
assert last_chunk2.output_text == "Hello!"
@pytest.mark.anyio
async def test_streaming_response_yields_expected_response_chunks(litellm_test_server):
api_base, _ = litellm_test_server
lm = dspy.LM(
model="openai/dspy-test-model",
api_base=api_base,
api_key="fakekey",
cache=False,
)
with dspy.context(lm=lm):
class TestSignature(dspy.Signature):
input_text: str = dspy.InputField()
output_text: str = dspy.OutputField()
program = dspy.streamify(dspy.Predict(TestSignature))
output_stream_from_program = streaming_response(program(input_text="Test"))
output_stream_for_server_response = streaming_response(output_stream_from_program)
output_chunks = [chunk async for chunk in output_stream_for_server_response]
assert all(chunk.startswith("data: ") for chunk in output_chunks)
assert 'data: {"prediction":{"output_text":"Hello!"}}\n\n' in output_chunks
assert output_chunks[-1] == "data: [DONE]\n\n"
@pytest.mark.anyio
async def test_default_status_streaming():
class MyProgram(dspy.Module):
def __init__(self):
self.generate_question = dspy.Tool(lambda x: f"What color is the {x}?", name="generate_question")
self.predict = dspy.Predict("question->answer")
def __call__(self, x: str):
question = self.generate_question(x=x)
return self.predict(question=question)
lm = dspy.utils.DummyLM([{"answer": "red"}, {"answer": "blue"}])
with dspy.context(lm=lm):
program = dspy.streamify(MyProgram())
output = program("sky")
status_messages = []
async for value in output:
if isinstance(value, StatusMessage):
status_messages.append(value)
assert len(status_messages) == 2
assert status_messages[0].message == "Calling tool generate_question..."
assert status_messages[1].message == "Tool calling finished! Querying the LLM with tool calling results..."
@pytest.mark.anyio
async def test_custom_status_streaming():
class MyProgram(dspy.Module):
def __init__(self):
self.generate_question = dspy.Tool(lambda x: f"What color is the {x}?", name="generate_question")
self.predict = dspy.Predict("question->answer")
def __call__(self, x: str):
question = self.generate_question(x=x)
return self.predict(question=question)
class MyStatusMessageProvider(StatusMessageProvider):
def tool_start_status_message(self, instance, inputs):
return "Tool starting!"
def tool_end_status_message(self, outputs):
return "Tool finished!"
def module_start_status_message(self, instance, inputs):
if isinstance(instance, dspy.Predict):
return "Predict starting!"
lm = dspy.utils.DummyLM([{"answer": "red"}, {"answer": "blue"}])
with dspy.context(lm=lm):
program = dspy.streamify(MyProgram(), status_message_provider=MyStatusMessageProvider())
output = program("sky")
status_messages = []
async for value in output:
if isinstance(value, StatusMessage):
status_messages.append(value)
assert len(status_messages) == 3
assert status_messages[0].message == "Tool starting!"
assert status_messages[1].message == "Tool finished!"
assert status_messages[2].message == "Predict starting!"
@pytest.mark.llm_call
@pytest.mark.anyio
async def test_stream_listener_chat_adapter(lm_for_test):
class MyProgram(dspy.Module):
def __init__(self):
self.predict1 = dspy.Predict("question->answer")
self.predict2 = dspy.Predict("question, answer->judgement")
def __call__(self, x: str, **kwargs):
answer = self.predict1(question=x, **kwargs)
judgement = self.predict2(question=x, answer=answer, **kwargs)
return judgement
my_program = MyProgram()
program = dspy.streamify(
my_program,
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer"),
dspy.streaming.StreamListener(signature_field_name="judgement"),
],
include_final_prediction_in_output_stream=False,
)
# Turn off the cache to ensure the stream is produced.
with dspy.context(lm=dspy.LM(lm_for_test, cache=False, temperature=0.0)):
output = program(x="why did a chicken cross the kitchen?")
all_chunks = []
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
assert all_chunks[0].predict_name == "predict1"
assert all_chunks[0].signature_field_name == "answer"
# The last chunk can be from either predictor because sometimes small LMs miss the `[[ ## completed ## ]]` marker,
# which results in an extra chunk that flushes out the buffer.
assert all_chunks[-2].predict_name == "predict2"
assert all_chunks[-2].signature_field_name == "judgement"
@pytest.mark.anyio
async def test_default_status_streaming_in_async_program():
class MyProgram(dspy.Module):
def __init__(self):
self.generate_question = dspy.Tool(lambda x: f"What color is the {x}?", name="generate_question")
self.predict = dspy.Predict("question->answer")
async def acall(self, x: str):
question = await self.generate_question.acall(x=x)
return await self.predict.acall(question=question)
lm = dspy.utils.DummyLM([{"answer": "red"}, {"answer": "blue"}])
with dspy.context(lm=lm):
program = dspy.streamify(MyProgram(), is_async_program=True)
output = program("sky")
status_messages = []
async for value in output:
if isinstance(value, StatusMessage):
status_messages.append(value)
assert len(status_messages) == 2
assert status_messages[0].message == "Calling tool generate_question..."
assert status_messages[1].message == "Tool calling finished! Querying the LLM with tool calling results..."
@pytest.mark.llm_call
@pytest.mark.anyio
async def test_stream_listener_json_adapter(lm_for_test):
class MyProgram(dspy.Module):
def __init__(self):
self.predict1 = dspy.Predict("question->answer")
self.predict2 = dspy.Predict("question, answer->judgement")
def __call__(self, x: str, **kwargs):
answer = self.predict1(question=x, **kwargs)
judgement = self.predict2(question=x, answer=answer, **kwargs)
return judgement
my_program = MyProgram()
program = dspy.streamify(
my_program,
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer"),
dspy.streaming.StreamListener(signature_field_name="judgement"),
],
include_final_prediction_in_output_stream=False,
)
# Turn off the cache to ensure the stream is produced.
with dspy.context(lm=dspy.LM(lm_for_test, cache=False, temperature=0.0), adapter=dspy.JSONAdapter()):
output = program(x="why did a chicken cross the kitchen?")
all_chunks = []
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
assert all_chunks[0].predict_name == "predict1"
assert all_chunks[0].signature_field_name == "answer"
assert all_chunks[0].is_last_chunk is False
assert all_chunks[-1].predict_name == "predict2"
assert all_chunks[-1].signature_field_name == "judgement"
@pytest.mark.anyio
async def test_streaming_handles_space_correctly():
my_program = dspy.Predict("question->answer")
program = dspy.streamify(
my_program, stream_listeners=[dspy.streaming.StreamListener(signature_field_name="answer")]
)
async def gpt_4o_mini_stream(*args, **kwargs):
yield ModelResponseStream(
model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="[[ ## answer ## ]]\n"))]
)
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="How "))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="are "))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="you "))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="doing?"))])
yield ModelResponseStream(
model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="\n\n[[ ## completed ## ]]"))]
)
with mock.patch("litellm.acompletion", side_effect=gpt_4o_mini_stream):
with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.ChatAdapter()):
output = program(question="What is the capital of France?")
all_chunks = []
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
assert "".join([chunk.chunk for chunk in all_chunks]) == "How are you doing?"
@pytest.mark.llm_call
def test_sync_streaming(lm_for_test):
class MyProgram(dspy.Module):
def __init__(self):
self.predict1 = dspy.Predict("question->answer")
self.predict2 = dspy.Predict("question, answer->judgement")
def __call__(self, x: str, **kwargs):
answer = self.predict1(question=x, **kwargs)
judgement = self.predict2(question=x, answer=answer, **kwargs)
return judgement
my_program = MyProgram()
program = dspy.streamify(
my_program,
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer"),
dspy.streaming.StreamListener(signature_field_name="judgement"),
],
include_final_prediction_in_output_stream=False,
async_streaming=False,
)
# Turn off the cache to ensure the stream is produced.
with dspy.context(lm=dspy.LM(lm_for_test, cache=False, temperature=0.0)):
output = program(x="why did a chicken cross the kitchen?")
all_chunks = []
for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
assert all_chunks[0].predict_name == "predict1"
assert all_chunks[0].signature_field_name == "answer"
assert all_chunks[0].is_last_chunk is False
# The last chunk can be from either predictor because sometimes small LMs miss the `[[ ## completed ## ]]` marker,
# which results in an extra chunk that flushes out the buffer.
assert all_chunks[-2].predict_name == "predict2"
assert all_chunks[-2].signature_field_name == "judgement"
def test_sync_status_streaming():
class MyProgram(dspy.Module):
def __init__(self):
self.generate_question = dspy.Tool(lambda x: f"What color is the {x}?", name="generate_question")
self.predict = dspy.Predict("question->answer")
def __call__(self, x: str):
question = self.generate_question(x=x)
return self.predict(question=question)
lm = dspy.utils.DummyLM([{"answer": "red"}, {"answer": "blue"}])
with dspy.context(lm=lm):
program = dspy.streamify(MyProgram())
output = program("sky")
sync_output = dspy.streaming.apply_sync_streaming(output)
status_messages = []
for value in sync_output:
if isinstance(value, StatusMessage):
status_messages.append(value)
assert len(status_messages) == 2
assert status_messages[0].message == "Calling tool generate_question..."
assert status_messages[1].message == "Tool calling finished! Querying the LLM with tool calling results..."
@pytest.mark.anyio
async def test_stream_listener_returns_correct_chunk_chat_adapter():
class MyProgram(dspy.Module):
def __init__(self):
super().__init__()
self.predict1 = dspy.Predict("question->answer")
self.predict2 = dspy.Predict("question, answer->judgement")
def forward(self, question, **kwargs):
answer = self.predict1(question=question, **kwargs).answer
judgement = self.predict2(question=question, answer=answer, **kwargs)
return judgement
async def gpt_4o_mini_stream_1(*args, **kwargs):
# Recorded streaming from openai/gpt-4o-mini
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="[["))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" answer"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ]]\n\n"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="To"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" get"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" to"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" the"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" other"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" side"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" of"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" the"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" dinner"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" plate"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="!\n\n[[ ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" completed"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ]]"))])
async def gpt_4o_mini_stream_2():
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="[[ ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" judgement"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ]]\n\n"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="The"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" answer"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" is"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" humorous"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" and"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" plays"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" on"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" the"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" classic"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" joke"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" format"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=".\n\n[[ ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" completed"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ]]"))])
stream_generators = [gpt_4o_mini_stream_1, gpt_4o_mini_stream_2]
async def completion_side_effect(*args, **kwargs):
return stream_generators.pop(0)() # return new async generator instance
with mock.patch("litellm.acompletion", side_effect=completion_side_effect):
program = dspy.streamify(
MyProgram(),
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer"),
dspy.streaming.StreamListener(signature_field_name="judgement"),
],
)
with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False)):
output = program(question="why did a chicken cross the kitchen?")
all_chunks = []
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
assert all_chunks[0].predict_name == "predict1"
assert all_chunks[0].signature_field_name == "answer"
assert all_chunks[0].chunk == "To"
assert all_chunks[1].chunk == " get"
assert all_chunks[2].chunk == " to"
assert all_chunks[3].chunk == " the"
assert all_chunks[4].chunk == " other"
assert all_chunks[5].chunk == " side"
assert all_chunks[6].chunk == " of"
assert all_chunks[7].chunk == " the"
assert all_chunks[8].chunk == " dinner"
assert all_chunks[9].chunk == " plate"
assert all_chunks[10].chunk == "!"
assert all_chunks[10].is_last_chunk is True
assert all_chunks[11].predict_name == "predict2"
assert all_chunks[11].signature_field_name == "judgement"
assert all_chunks[11].chunk == "The"
assert all_chunks[12].chunk == " answer"
assert all_chunks[13].chunk == " is"
assert all_chunks[14].chunk == " humorous"
assert all_chunks[15].chunk == " and"
assert all_chunks[16].chunk == " plays"
assert all_chunks[17].chunk == " on"
assert all_chunks[18].chunk == " the"
assert all_chunks[19].chunk == " classic"
assert all_chunks[20].chunk == " joke"
assert all_chunks[21].chunk == " format"
assert all_chunks[22].chunk == "."
assert all_chunks[22].is_last_chunk is True
@pytest.mark.anyio
async def test_stream_listener_returns_correct_chunk_json_adapter():
class MyProgram(dspy.Module):
def __init__(self):
super().__init__()
self.predict1 = dspy.Predict("question->answer")
self.predict2 = dspy.Predict("question,answer->judgement")
def forward(self, question, **kwargs):
answer = self.predict1(question=question, **kwargs).answer
judgement = self.predict2(question=question, answer=answer, **kwargs)
return judgement
async def gpt_4o_mini_stream_1(*args, **kwargs):
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content='{"'))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="answer"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content='":'))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="To"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" get"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" to"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" the"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" other"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" side"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" of"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" the"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" frying"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" pan"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content='!"'))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="}\n"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="None"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="None"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="None"))])
async def gpt_4o_mini_stream_2(*args, **kwargs):
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content='{"'))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="jud"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="gement"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content='":"'))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="The"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" answer"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" is"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" humorous"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" and"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" plays"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" on"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" the"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" very"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" funny"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" and"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" classic"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" joke"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" format"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content='."'))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="}"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="None"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="None"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="None"))])
with mock.patch(
"litellm.acompletion", new_callable=AsyncMock, side_effect=[gpt_4o_mini_stream_1(), gpt_4o_mini_stream_2()]
):
program = dspy.streamify(
MyProgram(),
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer"),
dspy.streaming.StreamListener(signature_field_name="judgement"),
],
)
with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.JSONAdapter()):
output = program(question="why did a chicken cross the kitchen?")
all_chunks = []
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
assert all_chunks[0].predict_name == "predict1"
assert all_chunks[0].signature_field_name == "answer"
assert all_chunks[0].chunk == "To"
assert all_chunks[1].chunk == " get"
assert all_chunks[2].chunk == " to"
assert all_chunks[3].chunk == " the"
assert all_chunks[4].chunk == " other"
assert all_chunks[5].chunk == " side"
assert all_chunks[6].chunk == " of"
assert all_chunks[7].chunk == " the"
assert all_chunks[8].chunk == " frying"
assert all_chunks[9].chunk == " pan"
assert all_chunks[10].chunk == "!"
assert all_chunks[10].is_last_chunk is True
assert all_chunks[11].predict_name == "predict2"
assert all_chunks[11].signature_field_name == "judgement"
assert all_chunks[11].chunk == "The"
assert all_chunks[12].chunk == " answer"
assert all_chunks[13].chunk == " is"
assert all_chunks[14].chunk == " humorous"
assert all_chunks[15].chunk == " and"
assert all_chunks[16].chunk == " plays"
assert all_chunks[17].chunk == " on"
assert all_chunks[18].chunk == " the"
assert all_chunks[19].chunk == " very"
assert all_chunks[20].chunk == " funny"
assert all_chunks[21].chunk == " and"
assert all_chunks[22].chunk == " classic"
assert all_chunks[23].chunk == " joke"
assert all_chunks[24].chunk == " format"
assert all_chunks[25].chunk == "."
assert all_chunks[25].is_last_chunk is True
@pytest.mark.anyio
async def test_stream_listener_returns_correct_chunk_chat_adapter_untokenized_stream():
class MyProgram(dspy.Module):
def __init__(self):
super().__init__()
self.predict1 = dspy.Predict("question->answer")
self.predict2 = dspy.Predict("question,answer->judgement")
def forward(self, question, **kwargs):
answer = self.predict1(question=question, **kwargs).answer
judgement = self.predict2(question=question, answer=answer, **kwargs)
return judgement
async def gemini_stream_1(*args, **kwargs):
yield ModelResponseStream(model="gemini", choices=[StreamingChoices(delta=Delta(content="[[ ##"))])
yield ModelResponseStream(model="gemini", choices=[StreamingChoices(delta=Delta(content=" answer ## ]]"))])
yield ModelResponseStream(
model="gemini", choices=[StreamingChoices(delta=Delta(content="To get to the other side."))]
)
yield ModelResponseStream(
model="gemini", choices=[StreamingChoices(delta=Delta(content="\n\n[[ ## completed ## ]]"))]
)
async def gemini_stream_2(*args, **kwargs):
yield ModelResponseStream(
model="gemini", choices=[StreamingChoices(delta=Delta(content="[[ ## judgement ## ]]\n\n"))]
)
yield ModelResponseStream(
model="gemini",
choices=[
StreamingChoices(
delta=Delta(
content=(
"The answer provides the standard punchline for this classic joke format, adapted to the "
"specific location mentioned in the question. It is the expected and appropriate response."
)
)
)
],
)
yield ModelResponseStream(
model="gemini",
choices=[StreamingChoices(delta=Delta(content="\n\n[[ ## completed ## ]]"))],
)
yield ModelResponseStream(model="gemini", choices=[StreamingChoices(delta=Delta(content="}\n"))])
with mock.patch("litellm.acompletion", new_callable=AsyncMock, side_effect=[gemini_stream_1(), gemini_stream_2()]):
program = dspy.streamify(
MyProgram(),
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer"),
dspy.streaming.StreamListener(signature_field_name="judgement"),
],
)
with dspy.context(lm=dspy.LM("gemini/gemini-2.5-flash", cache=False), adapter=dspy.ChatAdapter()):
output = program(question="why did a chicken cross the kitchen?")
all_chunks = []
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
assert all_chunks[0].predict_name == "predict1"
assert all_chunks[0].signature_field_name == "answer"
assert all_chunks[0].chunk == "To get to the other side."
assert all_chunks[1].predict_name == "predict2"
assert all_chunks[1].signature_field_name == "judgement"
assert all_chunks[1].chunk == (
"The answer provides the standard punchline for this classic joke format, adapted to the specific location "
"mentioned in the question. It is the expected and appropriate response."
)
@pytest.mark.anyio
async def test_stream_listener_missing_completion_marker_chat_adapter():
"""Test that streaming works correctly when LLM response omits a final completion marker.
This test verifies that:
1. All tokens are yielded including those in the buffer
2. The last chunk is properly marked with is_last_chunk=True
3. No tokens are lost when the completion marker is missing
"""
class MyProgram(dspy.Module):
def __init__(self):
super().__init__()
self.predict = dspy.Predict("question->answer")
def forward(self, question, **kwargs):
return self.predict(question=question, **kwargs)
async def incomplete_stream(*args, **kwargs):
"""Stream that includes start marker but MISSING completion marker"""
# Start marker
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="[[ ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" answer"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ## ]]\n\n"))])
# Content tokens - more than 10 to ensure buffering happens
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="This"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" is"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" a"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" test"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" response"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" with"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" many"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" tokens"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" to"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ensure"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" buffering"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" works"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" correctly"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="."))])
# NO COMPLETION MARKER
with mock.patch("litellm.acompletion", side_effect=incomplete_stream):
program = dspy.streamify(
MyProgram(),
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer"),
],
)
with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.ChatAdapter()):
output = program(question="Test question")
all_chunks = []
final_prediction = None
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
elif isinstance(value, dspy.Prediction):
final_prediction = value
full_content = "".join([chunk.chunk for chunk in all_chunks])
expected_content = "This is a test response with many tokens to ensure buffering works correctly."
assert full_content == expected_content
assert final_prediction.answer == expected_content
@pytest.mark.anyio
async def test_stream_listener_returns_correct_chunk_json_adapter_untokenized_stream():
class MyProgram(dspy.Module):
def __init__(self):
super().__init__()
self.predict1 = dspy.Predict("question->answer")
self.predict2 = dspy.Predict("question,answer->judgement")
def forward(self, question, **kwargs):
answer = self.predict1(question=question, **kwargs).answer
judgement = self.predict2(question=question, answer=answer, **kwargs)
return judgement
async def gemini_stream_1(*args, **kwargs):
yield ModelResponseStream(model="gemini", choices=[StreamingChoices(delta=Delta(content="{\n"))])
yield ModelResponseStream(
model="gemini", choices=[StreamingChoices(delta=Delta(content=' "answer": "To get to'))]
)
yield ModelResponseStream(
model="gemini", choices=[StreamingChoices(delta=Delta(content=' the other side... of the cutting board!"'))]
)
yield ModelResponseStream(model="gemini", choices=[StreamingChoices(delta=Delta(content="}\n"))])
async def gemini_stream_2(*args, **kwargs):
yield ModelResponseStream(model="gemini", choices=[StreamingChoices(delta=Delta(content="{\n"))])
yield ModelResponseStream(
model="gemini", choices=[StreamingChoices(delta=Delta(content=' "judgement": "The'))]
)
yield ModelResponseStream(
model="gemini",
choices=[
StreamingChoices(
delta=Delta(
content=' answer provides a humorous and relevant punchline to the classic joke setup."'
)
)
],
)
yield ModelResponseStream(model="gemini", choices=[StreamingChoices(delta=Delta(content="}\n"))])
with mock.patch("litellm.acompletion", new_callable=AsyncMock, side_effect=[gemini_stream_1(), gemini_stream_2()]):
program = dspy.streamify(
MyProgram(),
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer"),
dspy.streaming.StreamListener(signature_field_name="judgement"),
],
)
with dspy.context(lm=dspy.LM("gemini/gemini-2.5-flash", cache=False), adapter=dspy.JSONAdapter()):
output = program(question="why did a chicken cross the kitchen?")
all_chunks = []
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
assert all_chunks[0].predict_name == "predict1"
assert all_chunks[0].signature_field_name == "answer"
assert all_chunks[0].chunk == "To get to the other side... of the cutting board!"
assert all_chunks[1].predict_name == "predict2"
assert all_chunks[1].signature_field_name == "judgement"
assert all_chunks[1].chunk == "The answer provides a humorous and relevant punchline to the classic joke setup."
@pytest.mark.anyio
async def test_status_message_non_blocking():
def dummy_tool():
time.sleep(1)
return "dummy_tool_output"
class MyProgram(dspy.Module):
def forward(self, question, **kwargs):
dspy.Tool(dummy_tool)()
return dspy.Prediction(answer="dummy_tool_output")
program = dspy.streamify(MyProgram(), status_message_provider=StatusMessageProvider())
with mock.patch("litellm.acompletion", new_callable=AsyncMock, side_effect=[dummy_tool]):
with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False)):
output = program(question="why did a chicken cross the kitchen?")
timestamps = []
async for value in output:
if isinstance(value, dspy.streaming.StatusMessage):
timestamps.append(time.time())
# timestamps[0]: tool start message
# timestamps[1]: tool end message
# There should be ~1 second delay between the tool start and end messages because we explicitly sleep for 1 second
# in the tool.
assert timestamps[1] - timestamps[0] >= 1
@pytest.mark.anyio
async def test_status_message_non_blocking_async_program():
async def dummy_tool():
await asyncio.sleep(1)
return "dummy_tool_output"
class MyProgram(dspy.Module):
async def aforward(self, question, **kwargs):
await dspy.Tool(dummy_tool).acall()
return dspy.Prediction(answer="dummy_tool_output")
program = dspy.streamify(MyProgram(), status_message_provider=StatusMessageProvider(), is_async_program=True)
with mock.patch("litellm.acompletion", new_callable=AsyncMock, side_effect=[dummy_tool]):
with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False)):
output = program(question="why did a chicken cross the kitchen?")
timestamps = []
async for value in output:
if isinstance(value, dspy.streaming.StatusMessage):
timestamps.append(time.time())
# timestamps[0]: tool start message
# timestamps[1]: tool end message
# There should be ~1 second delay between the tool start and end messages because we explicitly sleep for 1 second
# in the tool.
assert timestamps[1] - timestamps[0] >= 1
@pytest.mark.anyio
async def test_stream_listener_allow_reuse():
class MyProgram(dspy.Module):
def __init__(self):
super().__init__()
self.predict = dspy.Predict("question->answer")
def forward(self, question, **kwargs):
self.predict(question=question, **kwargs)
return self.predict(question=question, **kwargs)
program = dspy.streamify(
MyProgram(),
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer", allow_reuse=True),
],
)
async def gpt_4o_mini_stream(*args, **kwargs):
# Recorded streaming from openai/gpt-4o-mini
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="[["))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" answer"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ]]\n\n"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="To"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" get"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" to"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" the"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" other"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" side"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="!\n\n[[ ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" completed"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ]]"))])
stream_generators = [gpt_4o_mini_stream, gpt_4o_mini_stream]
async def completion_side_effect(*args, **kwargs):
return stream_generators.pop(0)() # return new async generator instance
with mock.patch("litellm.acompletion", side_effect=completion_side_effect):
with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False)):
output = program(question="why did a chicken cross the kitchen?")
all_chunks = []
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
concat_message = "".join([chunk.chunk for chunk in all_chunks])
# The listener functions twice.
assert concat_message == "To get to the other side!To get to the other side!"
@pytest.mark.anyio
async def test_stream_listener_returns_correct_chunk_xml_adapter():
class MyProgram(dspy.Module):
def __init__(self):
super().__init__()
self.predict1 = dspy.Predict("question->answer")
self.predict2 = dspy.Predict("question,answer->judgement")
def forward(self, question, **kwargs):
answer = self.predict1(question=question, **kwargs).answer
judgement = self.predict2(question=question, answer=answer, **kwargs)
return judgement
async def xml_stream_1(*args, **kwargs):
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="<"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="answer"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=">"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="To"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" get"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" to"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" the"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" other"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" side"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="!"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="<"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="/answer"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=">"))])
async def xml_stream_2(*args, **kwargs):
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="<"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="judgement"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=">"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="The"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" answer"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" is"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" humorous"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="."))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="<"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="/judgement"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=">"))])
stream_generators = [xml_stream_1, xml_stream_2]
async def completion_side_effect(*args, **kwargs):
return stream_generators.pop(0)()
with mock.patch("litellm.acompletion", side_effect=completion_side_effect):
program = dspy.streamify(
MyProgram(),
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer"),
dspy.streaming.StreamListener(signature_field_name="judgement"),
],
)
with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.XMLAdapter()):
output = program(question="why did a chicken cross the kitchen?")
all_chunks = []
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
# Verify answer chunks
answer_chunks = [chunk for chunk in all_chunks if chunk.signature_field_name == "answer"]
assert len(answer_chunks) > 0
assert answer_chunks[0].predict_name == "predict1"
assert "".join([chunk.chunk for chunk in answer_chunks]) == "To get to the other side!"
# Verify judgement chunks
judgement_chunks = [chunk for chunk in all_chunks if chunk.signature_field_name == "judgement"]
assert len(judgement_chunks) > 0
assert judgement_chunks[0].predict_name == "predict2"
assert "".join([chunk.chunk for chunk in judgement_chunks]) == "The answer is humorous."
@pytest.mark.anyio
async def test_streaming_allows_custom_chunk_types():
@dataclass
class CustomChunk:
text: str
class MyProgram(dspy.Module):
def forward(self, question, **kwargs):
async def send_to_stream():
chunk = CustomChunk(text="hello")
await dspy.settings.send_stream.send(chunk)
syncified_send_to_stream = syncify(send_to_stream)
syncified_send_to_stream()
return dspy.Prediction(answer="dummy output")
program = dspy.streamify(MyProgram())
output = program(question="why did a chicken cross the kitchen?")
all_chunks = []
async for value in output:
all_chunks.append(value)
assert isinstance(all_chunks[0], CustomChunk)
assert isinstance(all_chunks[1], dspy.Prediction)
@pytest.mark.anyio
async def test_streaming_allows_custom_streamable_type():
class CustomType(Type):
message: str
@classmethod
def is_streamable(cls) -> bool:
return True
@classmethod
def parse_stream_chunk(cls, chunk):
return CustomType(message=chunk.choices[0].delta.content)
@classmethod
def parse_lm_response(cls, response: dict) -> "CustomType":
return CustomType(message=response.split("\n\n")[0])
class CustomSignature(dspy.Signature):
question: str = dspy.InputField()
answer: CustomType = dspy.OutputField()
program = dspy.streamify(
dspy.Predict(CustomSignature),
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer"),
],
)
async def stream(*args, **kwargs):
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="Hello"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="World"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="\n\n"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="[[ ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" completed"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ]]"))])
with mock.patch("litellm.acompletion", side_effect=stream):
with dspy.context(
lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.ChatAdapter(native_response_types=[CustomType])
):
output = program(question="why did a chicken cross the kitchen?")
all_chunks = []
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse):
all_chunks.append(value)
elif isinstance(value, dspy.Prediction):
assert isinstance(value.answer, CustomType)
assert value.answer.message == "HelloWorld"
assert all(isinstance(chunk.chunk, CustomType) for chunk in all_chunks)
@pytest.mark.anyio
async def test_streaming_with_citations():
class AnswerWithSources(dspy.Signature):
"""Answer questions using provided documents with citations."""
documents: list[Document] = dspy.InputField()
question: str = dspy.InputField()
answer: str = dspy.OutputField()
citations: Citations = dspy.OutputField()
class MyProgram(dspy.Module):
def __init__(self):
super().__init__()
self.predict = dspy.Predict(AnswerWithSources)
def forward(self, documents, question, **kwargs):
return self.predict(documents=documents, question=question, **kwargs)
async def citation_stream(*args, **kwargs):
# Stream chunks with citation data in provider_specific_fields
# To verify the realistic scenario with more than 10 chunks in the stream, include more than 10 chunks before the citation.
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content="[[ ##"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content=" answer"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content=" ## ]]\n\n"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content="A"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content="c"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content="c"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content="o"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content="r"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content="d"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content="i"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content="n"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content="g"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content=" to "))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content="the references,"))])
yield ModelResponseStream(
model="claude",
choices=[
StreamingChoices(
delta=Delta(
content="",
provider_specific_fields={
"citation": {
"type": "char_location",
"cited_text": "water boils at 100°C",
"document_index": 0,
"document_title": "Physics Facts",
"start_char_index": 0,
"end_char_index": 19,
}
},
)
)
],
)
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content=" water"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content=" boils"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content=" at"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content=" 100°C"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content=".\n\n[[ ##"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content=" completed"))])
yield ModelResponseStream(model="claude", choices=[StreamingChoices(delta=Delta(content=" ## ]]"))])
# Mock the final response choice to include provider_specific_fields with citations
with mock.patch("litellm.acompletion", return_value=citation_stream()):
program = dspy.streamify(
MyProgram(),
stream_listeners=[
dspy.streaming.StreamListener(signature_field_name="answer"),
dspy.streaming.StreamListener(signature_field_name="citations"),
],
)
# Create test documents
docs = [Document(data="Water boils at 100°C at standard pressure.", title="Physics Facts")]
with dspy.context(
lm=dspy.LM("anthropic/claude-3-5-sonnet-20241022", cache=False),
adapter=dspy.ChatAdapter(native_response_types=[Citations]),
):
output = program(documents=docs, question="What temperature does water boil?")
citation_chunks = []
answer_chunks = []
final_prediction = None
async for value in output:
if isinstance(value, dspy.streaming.StreamResponse) and value.signature_field_name == "citations":
citation_chunks.append(value)
elif isinstance(value, dspy.streaming.StreamResponse) and value.signature_field_name == "answer":
answer_chunks.append(value.chunk)
elif isinstance(value, dspy.Prediction):
final_prediction = value
# Test that we received citation chunks from streaming
assert len(citation_chunks) > 0
citation_chunk = citation_chunks[0]
assert isinstance(citation_chunk.chunk, Citations)
assert len(citation_chunk.chunk) == 1
assert citation_chunk.chunk[0].cited_text == "water boils at 100°C"
assert citation_chunk.chunk[0].document_title == "Physics Facts"
# Verify the answer chunks are correct
assert "".join(answer_chunks) == "According to the references, water boils at 100°C."
# Test that prediction contains the expected fields
assert final_prediction is not None
assert hasattr(final_prediction, "answer")
assert hasattr(final_prediction, "citations")
assert final_prediction.answer == "According to the references, water boils at 100°C."
def test_stream_listener_could_form_end_identifier_chat_adapter():
listener = dspy.streaming.StreamListener(signature_field_name="answer")
# Should return True for partial bracket sequences
assert listener._could_form_end_identifier("some text [", "ChatAdapter") is True
assert listener._could_form_end_identifier("some text [[", "ChatAdapter") is True
assert listener._could_form_end_identifier("some text [[ ", "ChatAdapter") is True
assert listener._could_form_end_identifier("some text [[ #", "ChatAdapter") is True
assert listener._could_form_end_identifier("some text [[ ##", "ChatAdapter") is True
# Should return True for partial field names after "[[ ##"
assert listener._could_form_end_identifier("some text [[ ## com", "ChatAdapter") is True
assert listener._could_form_end_identifier("some text [[ ## completed", "ChatAdapter") is True
# Should return False for text that clearly cannot form the pattern
assert listener._could_form_end_identifier("hello world", "ChatAdapter") is False
assert listener._could_form_end_identifier("some text", "ChatAdapter") is False
assert listener._could_form_end_identifier("answer: hello", "ChatAdapter") is False
def test_stream_listener_could_form_end_identifier_json_adapter():
listener = dspy.streaming.StreamListener(signature_field_name="output")
# Should return True for partial quote/brace sequences
assert listener._could_form_end_identifier('some text "', "JSONAdapter") is True
assert listener._could_form_end_identifier('some text ",', "JSONAdapter") is True
assert listener._could_form_end_identifier('some text " ', "JSONAdapter") is True
assert listener._could_form_end_identifier('some text "}', "JSONAdapter") is True
# Should return False for text that cannot form the pattern
assert listener._could_form_end_identifier("hello world", "JSONAdapter") is False
assert listener._could_form_end_identifier("some text", "JSONAdapter") is False
def test_stream_listener_could_form_end_identifier_xml_adapter():
listener = dspy.streaming.StreamListener(signature_field_name="result")
# Should return True for partial closing tag
assert listener._could_form_end_identifier("some text <", "XMLAdapter") is True
assert listener._could_form_end_identifier("some text </", "XMLAdapter") is True
assert listener._could_form_end_identifier("some text </result", "XMLAdapter") is True
# Should return False for text that cannot form the pattern
assert listener._could_form_end_identifier("hello world", "XMLAdapter") is False
assert listener._could_form_end_identifier("some text", "XMLAdapter") is False
```