[ { "path": ".github/workflows/build.yml", "content": "name: Build\non:\n pull_request:\n branches: [main]\n push:\n branches: [develop]\n\njobs:\n check:\n name: Check\n runs-on: ubuntu-latest\n steps:\n - name: Checkout sources\n uses: actions/checkout@v2\n\n - name: Install stable toolchain\n uses: actions-rs/toolchain@v1\n with:\n profile: minimal\n toolchain: stable\n override: true\n\n - uses: Swatinem/rust-cache@v1\n\n - name: Run cargo check\n uses: actions-rs/cargo@v1\n with:\n command: check\n\n test:\n name: Test Suite\n strategy:\n matrix:\n # TODO: #7 Add tests for Windows & macOS.\n # os: [ubuntu-latest, macos-latest, windows-latest]\n os: [ubuntu-latest]\n rust: [stable]\n runs-on: ${{ matrix.os }}\n steps:\n - name: Checkout sources\n uses: actions/checkout@v2\n\n - name: Install stable toolchain\n uses: actions-rs/toolchain@v1\n with:\n profile: minimal\n toolchain: ${{ matrix.rust }}\n override: true\n\n - uses: Swatinem/rust-cache@v1\n\n - name: Unit tests\n uses: actions-rs/cargo@v1\n with:\n command: test\n\n - name: Doc tests\n uses: actions-rs/cargo@v1\n with:\n command: test\n args: --doc\n\n # TODO: Test with Ollama.\n # - name: Install and run ollama\n # run: |\n # ./scripts/ci.sh\n\n # - name: Examples (Ollama)\n # uses: actions-rs/cargo@v1\n # with:\n # command: test\n # args: -p core --examples -- ollama\n\n - name: Examples (OpenAI)\n run: ./scripts/examples.sh openai\n env:\n OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY_E2E }}\n\n - name: Examples (Anthropic)\n run: ./scripts/examples.sh anthropic\n env:\n ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY_E2E }}\n\n lints:\n name: Lints\n runs-on: ubuntu-latest\n steps:\n - name: Checkout sources\n uses: actions/checkout@v2\n with:\n submodules: true\n\n - name: Install stable toolchain\n uses: actions-rs/toolchain@v1\n with:\n profile: minimal\n toolchain: stable\n override: true\n components: rustfmt, clippy\n\n - uses: Swatinem/rust-cache@v1\n\n - name: Run cargo fmt\n uses: actions-rs/cargo@v1\n with:\n command: fmt\n args: --all -- --check\n\n - name: Run cargo clippy\n uses: actions-rs/cargo@v1\n with:\n command: clippy\n args: -- -D warnings\n\n # - name: Run rustdoc lints\n # uses: actions-rs/cargo@v1\n # env:\n # RUSTDOCFLAGS: \"-D missing_docs -D rustdoc::missing_doc_code_examples\"\n # with:\n # command: doc\n # args: --workspace --all-features --no-deps --document-private-items\n" }, { "path": ".github/workflows/release.yml", "content": "name: Release\non:\n push:\n tags: [\"[0-9]+.[0-9]+.[0-9]+\"]\n\nenv:\n BRANCH_NAME: ${{ github.head_ref || github.ref_name }} \n\njobs:\n release:\n name: Release\n runs-on: ubuntu-latest\n steps:\n - name: Exit if not on main branch\n if: endsWith(github.event.base_ref, 'main') == false\n run: exit 1\n\n - name: Checkout sources\n uses: actions/checkout@v2\n\n - name: Install stable toolchain\n uses: actions-rs/toolchain@v1\n with:\n profile: minimal\n toolchain: stable\n override: true\n\n - name: Publish `orch_response`\n run: cargo publish -p orch_response --token ${{ secrets.CRATES_IO_API_TOKEN }}\n \n - name: Publish `orch_response_derive`\n run: cargo publish -p orch_response_derive --token ${{ secrets.CRATES_IO_API_TOKEN }}\n\n - name: Publish `orch`\n run: cargo publish -p orch --token ${{ secrets.CRATES_IO_API_TOKEN }}\n \n" }, { "path": ".gitignore", "content": "# Generated by Cargo\n# will have compiled files and executables\ndebug/\ntarget/\n\n# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries\n# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html\nCargo.lock\n\n# These are backup files generated by rustfmt\n**/*.rs.bk\n\n# MSVC Windows builds of rustc generate these, which store debugging information\n*.pdb\n\n.env*\n\n.DS_Store" }, { "path": ".rusfmt.toml", "content": "max_width = 140\n" }, { "path": ".vscode/settings.json", "content": "{\n\t\"rust-analyzer.linkedProjects\": [\n\t\t\"./Cargo.toml\"\n\t]\n}" }, { "path": "CHANGELOG.md", "content": "# Changelog\n\n# Version 0.0.16\n\n- Added ability to configure the OpenAI API endpoint\n- Added support for Anthropic models\n- Fixed issue with `is_local` producing the incorrect result\n- Improved examples for the `orch` crate\n\n# Version 0.0.15\n\n- Removed dimensions from the OpenAI embedding model (no such requirements, compared to Ollama)\n\n# Version 0.0.14\n\n- Added `LangaugeModelProvider::is_local` method\n\n# Version 0.0.13\n\nNo functional changes.\n\n# Version 0.0.12\n\n- Fixed the proc macro for the variants of the response options ([PR #9](https://github.com/guywaldman/orch/pull/9))\n- Added support for \"alignment\", improved documentation, added examples ([PR #11](https://github.com/guywaldman/orch/pull/11))\n\n# Version 0.0.11\n\nNo functional changes.\n\n# Version 0.0.10\n\nNo functional changes.\n\n# Version 0.0.9\n\nNo functional changes.\n\n# Version 0.0.8\n\n- Added support for boolean fields in the response options\n\n# Version 0.0.7\n\n- Fixed issue where the `orch` crate was not used for types in the proc macros\n- Fixed issue where multiple fields in a response option would fail the proc macro\n\n# Version 0.0.6\n\nNo functional changes.\n\n# Version 0.0.5\n\n- Fixed an issue where the proc macros were not exposed directly from `orch`\n\n# Version 0.0.4\n\nNo functional changes.\n\n# Version 0.0.3\n\n- Added support for streaming responses\n- Added support for structured data generation\n- Added a convenience proc macro (`#[derive(OrchResponseOptions)]`) for generating structured data generation options\n- Added support for Open AI\n\n## Version 0.0.2\n\nNo functional changes.\n\n## Version 0.0.1\n\nInitial release.\n" }, { "path": "CONTRIBUTING.md", "content": "# Contributing to `orch`\n\nThank you for your interest in contributing to `orch`!\n\nPlease follow this process for contributing code:\n\n1. Open an issue to discuss the changes you want to make\n1. Wait for maintainers to approve this change\n > This step is solely to reduce noise and avoid redundant work for implementing a change that isn't accepted. The assumption is that that a maintainer will review the issue in reasonable time.\n\n## Contribution Workflow\n\n1. Fork the respository\n1. Implement the change you wish to make\n1. Open a pull request that references the relevant issue\n1. Make sure that the CI passes\n1. Wait for maintainers to review and approve the pull request\n\n## Development Environment\n\n### Prerequisites\n\n1. Rust toolchain and Cargo (MSVC: 1.78.0)\n\n### Workflow\n\n```shell\ngit clone https://github.com/{your-username}/magic-cli\ncd magic-cli\n\n# ...Implement the changes you wish to make\n\ncargo fmt\ncargo test\n\n# ...Commit your changes\n\n# ...Push your changes\n\n# ...Open a pull request\n```\n" }, { "path": "Cargo.toml", "content": "[workspace]\nresolver = \"2\"\n\nmembers = [\"core\", \"response\", \"response_derive\"]\n" }, { "path": "LICENSE.md", "content": "MIT License\n\nCopyright (c) [year] [fullname]\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "README.md", "content": "# orch\n\n![Crates.io Version](https://img.shields.io/crates/v/orch?link=https%3A%2F%2Fcrates.io%2Fcrates%2Forch)\n![Crates.io Total Downloads](https://img.shields.io/crates/d/orch?link=https%3A%2F%2Fcrates.io%2Fcrates%2Forch)\n\n`orch` is a library for building language model powered applications and agents for the Rust programming language.\nIt was primarily built for usage in [magic-cli](https://github.com/guywaldman/magic-cli), but can be used in other contexts as well.\n\n> [!NOTE]\n>\n> If the project gains traction, this can be compiled as an addon to other languages such as Python or a standalone WebAssembly module.\n\n# Installation\n\n```shell\ncargo add orch\ncargo add orch_response\n```\n\nAlternatively, add `orch` as a dependency to your `Cargo.toml` file:\n\n```toml\n[dependencies]\norch = \"*\" # Substitute with the latest version\norch_response = \"*\" # Substitute with the latest version\n```\n\n# Basic Usage\n\n## Simple Text Generation\n\n```rust no_run\nuse orch::execution::*;\nuse orch::lm::*;\n\n#[tokio::main]\nasync fn main() {\n let lm = OllamaBuilder::new().try_build().unwrap();\n let executor = TextExecutorBuilder::new().with_lm(&lm).try_build().unwrap();\n let response = executor.execute(\"What is 2+2?\").await.expect(\"Execution failed\");\n println!(\"{}\", response.content);\n}\n```\n\n## Streaming Text Generation\n\n```rust no_run\nuse orch::execution::*;\nuse orch::lm::*;\nuse tokio_stream::StreamExt;\n\n#[tokio::main]\nasync fn main() {\n let lm = OllamaBuilder::new().try_build().unwrap();\n let executor = TextExecutorBuilder::new().with_lm(&lm).try_build().unwrap();\n let mut response = executor.execute_stream(\"What is 2+2?\").await.expect(\"Execution failed\");\n while let Some(chunk) = response.stream.next().await {\n match chunk {\n Ok(chunk) => print!(\"{chunk}\"),\n Err(e) => {\n println!(\"Error: {e}\");\n break;\n }\n }\n }\n println!();\n}\n```\n\n## Structured Data Generation\n\n```rust no_run\nuse orch::execution::*;\nuse orch::lm::*;\nuse orch::response::*;\n\n#[derive(Variants, serde::Deserialize)]\npub enum ResponseVariants {\n Answer(AnswerResponseVariant),\n Fail(FailResponseVariant),\n}\n\n#[derive(Variant, serde::Deserialize)]\n#[variant(\n variant = \"Answer\",\n scenario = \"You know the capital city of the country\",\n description = \"Capital city of the country\"\n)]\npub struct AnswerResponseVariant {\n #[schema(\n description = \"Capital city of the received country\",\n example = \"London\"\n )]\n pub capital: String,\n}\n\n#[derive(Variant, serde::Deserialize)]\n#[variant(\n variant = \"Fail\",\n scenario = \"You don't know the capital city of the country\",\n description = \"Reason why the capital city is not known\"\n)]\npub struct FailResponseVariant {\n #[schema(\n description = \"Reason why the capital city is not known\",\n example = \"Country 'foobar' does not exist\"\n )]\n pub reason: String,\n}\n\n#[tokio::main]\nasync fn main() {\n let lm = OllamaBuilder::new().try_build().unwrap();\n let executor = StructuredExecutorBuilder::new()\n .with_lm(&lm)\n .with_preamble(\"You are a geography expert who helps users understand the capital city of countries around the world.\")\n\t\t.with_options(Box::new(variants!(ResponseVariants)))\n .try_build()\n .unwrap();\n let response = executor\n .execute(\"What is the capital of Fooland?\")\n .await\n .expect(\"Execution failed\");\n\n println!(\"Response:\");\n match response.content {\n ResponseVariants::Answer(answer) => {\n println!(\"Capital city: {}\", answer.capital);\n }\n ResponseVariants::Fail(fail) => {\n println!(\"Model failed to generate a response: {}\", fail.reason);\n }\n }\n}\n```\n\n## Embedding Generation\n\n```rust no_run\nuse orch::execution::*;\nuse orch::lm::*;\n\n#[tokio::main]\nasync fn main() {\n let lm = OllamaBuilder::new().try_build().unwrap();\n let executor = TextExecutorBuilder::new()\n .with_lm(&lm)\n .try_build()\n .unwrap();\n let embedding = executor\n .generate_embedding(\"Phrase to generate an embedding for\")\n .await\n .expect(\"Execution failed\");\n\n println!(\"Embedding:\");\n println!(\"{:?}\", embedding);\n}\n```\n\n## More Examples\n\nSee the [examples](https://github.com/guywaldman/orch/tree/main/core/examples) directory for usage examples.\n" }, { "path": "RELEASE.md", "content": "# Release Process\n\n1. Update the versions in all package `Cargo.toml` files and version in the README.md file\n > Alternatively, rename all versions that are not in [CHANGELOG.md](CHANGELOG.md) to the next version.\n1. Update [CHANGELOG.md](CHANGELOG.md)\n1. Push the changes to the `main` branch with a tag (e.g., `0.0.6`)\n" }, { "path": "core/Cargo.toml", "content": "[package]\nname = \"orch\"\nversion = \"0.0.16\"\nedition = \"2021\"\nlicense = \"MIT\"\ndescription = \"Language model orchestration library\"\nhomepage = \"https://github.com/guywaldman/orch\"\nrepository = \"https://github.com/guywaldman/orch\"\nkeywords = [\"llm\", \"openai\", \"ollama\", \"rust\"]\n\n[dependencies]\norch_response = { path = \"../response\", version = \"0.0.16\" }\norch_response_derive = { path = \"../response_derive\", version = \"0.0.16\" }\nasync-gen = \"0.2.3\"\ndotenv = \"0.15.0\"\nopenai-api-rs = \"5.0.2\"\nreqwest = { version = \"0.12.5\", features = [\"blocking\"] }\nserde = { version = \"1.0.164\", features = [\"derive\"] }\nserde_json = \"1.0.97\"\nthiserror = \"1.0.63\"\ntokio = { version = \"1.28.2\", features = [\"rt\", \"macros\"] }\ntokio-stream = \"0.1.15\"\nasync-trait = \"0.1.81\"\ndyn-clone = \"1.0.17\"\nasync-recursion = \"1.1.1\"\n" }, { "path": "core/README.md", "content": "# orch\n\nSee the [main README](../README.md) for more information.\n" }, { "path": "core/examples/alignment.rs", "content": "#![allow(dead_code)]\n\nuse orch::alignment::AlignmentStrategyBuilder;\nuse orch::execution::*;\nuse orch::response::*;\n\nmod example_utils;\nuse example_utils::get_lm;\n\n#[derive(Variants, serde::Deserialize)]\npub enum ResponseVariants {\n Answer(AnswerResponseVariant),\n Fail(FailResponseVariant),\n}\n\n#[derive(Variant, serde::Deserialize)]\n#[variant(\n variant = \"Answer\",\n scenario = \"You know the answer\",\n description = \"Result of the calculation\"\n)]\npub struct AnswerResponseVariant {\n #[schema(description = \"Result of the calculation\", example = \"42\")]\n pub result: String,\n}\n\n#[derive(Variant, serde::Deserialize)]\n#[variant(\n variant = \"Fail\",\n scenario = \"You don't know the answer\",\n description = \"Reason why the answer is not known\"\n)]\npub struct FailResponseVariant {\n #[schema(\n description = \"Reason why the answer is not known\",\n example = \"The phrase is not a mathematical related expression\"\n )]\n pub reason: String,\n}\n\n#[tokio::main]\nasync fn main() {\n let (lm, _) = get_lm();\n // In this example, we use the same LLM for the correction as for the main task.\n // This could be replaced by a smaller LM.\n let (corrector_lm, _) = get_lm();\n\n // We define an alignment strategy that uses the correction model.\n let alignment_strategy = AlignmentStrategyBuilder::new()\n .with_retries(2)\n .with_lm(&*corrector_lm)\n .try_build()\n .unwrap();\n\n let executor = StructuredExecutorBuilder::new()\n\t.with_lm(&*lm)\n\t.with_preamble(\"\n\t\tYou are a mathematician who helps users understand the result of mathematical expressions.\n\t\tYou will receive a mathematical expression, and you will need to provide the result of that expression.\n\t\")\n\t.with_options(Box::new(variants!(ResponseVariants)))\n .with_alignment(alignment_strategy)\n\t.try_build()\n\t.unwrap();\n let response = executor.execute(\"2 + 2\").await.expect(\"Execution failed\");\n\n match response.content {\n ResponseVariants::Answer(answer) => {\n println!(\"Result: {}\", answer.result);\n }\n ResponseVariants::Fail(fail) => {\n println!(\"Model failed to generate a response: {}\", fail.reason);\n }\n }\n}\n" }, { "path": "core/examples/embeddings.rs", "content": "//! This example demonstrates how to use the `Executor` to generate embeddings from the language model.\n//! Run like so: `cargo run --example embeddings`\n\nmod example_utils;\nuse example_utils::get_lm;\n\nuse orch::{execution::*, lm::LanguageModelProvider};\n\n#[tokio::main]\nasync fn main() {\n let (lm, provider) = get_lm();\n\n if provider == LanguageModelProvider::Anthropic {\n println!(\"Anthropic does not have built-in embedding models. Skipping example.\");\n return;\n }\n\n let text = \"Lorem ipsum\";\n\n println!(\"Text: {text}\");\n println!(\"---\");\n\n let executor = TextExecutorBuilder::new()\n .with_lm(&*lm)\n .try_build()\n .unwrap();\n let embedding = executor\n .generate_embedding(text)\n .await\n .expect(\"Execution failed\");\n\n println!(\"Embedding:\");\n println!(\"{:?}\", embedding);\n}\n" }, { "path": "core/examples/example_utils.rs", "content": "use orch::lm::{\n AnthropicBuilder, LanguageModel, LanguageModelBuilder, LanguageModelProvider, OllamaBuilder,\n OpenAiBuilder,\n};\n\npub fn get_lm() -> (Box, LanguageModelProvider) {\n let args = std::env::args().collect::>();\n let provider_name = args.get(1).unwrap_or_else(|| {\n eprintln!(\"ERROR: Please provide a provider name\");\n std::process::exit(1);\n });\n let provider = LanguageModelProvider::try_from(provider_name.as_str())\n .expect(\"Invalid provider name. Supported values: 'ollama', 'openai', 'anthropic'\");\n\n let open_ai_api_key = {\n if provider == LanguageModelProvider::OpenAi {\n std::env::var(\"OPENAI_API_KEY\")\n .unwrap_or_else(|_| panic!(\"OPENAI_API_KEY environment variable not set\"))\n } else {\n String::new()\n }\n };\n let anthropic_api_key = {\n if provider == LanguageModelProvider::Anthropic {\n std::env::var(\"ANTHROPIC_API_KEY\")\n .unwrap_or_else(|_| panic!(\"ANTHROPIC_API_KEY environment variable not set\"))\n } else {\n String::new()\n }\n };\n let lm: Box = match provider {\n LanguageModelProvider::Ollama => Box::new(OllamaBuilder::new().try_build().unwrap()),\n LanguageModelProvider::OpenAi => Box::new(\n OpenAiBuilder::new()\n .with_api_key(open_ai_api_key)\n .try_build()\n .unwrap(),\n ),\n LanguageModelProvider::Anthropic => Box::new(\n AnthropicBuilder::new()\n .with_api_key(anthropic_api_key)\n .try_build()\n .unwrap(),\n ),\n };\n\n (lm, provider)\n}\n\n#[allow(dead_code)]\nfn main() {}\n" }, { "path": "core/examples/structured_data_generation_blog.rs", "content": "//! This example demonstrates how to use the `Executor` to generate a structured response from the LLM.\n//! Run like so: `cargo run --example structured_data_generation_blog -- blog.md`\n\n#![allow(dead_code)]\n\nuse orch::alignment::AlignmentStrategyBuilder;\nuse orch::execution::*;\nuse orch::response::*;\n\nmod example_utils;\nuse example_utils::get_lm;\n\n#[derive(Variants, serde::Deserialize)]\n#[serde(tag = \"response_type\")]\npub enum ResponseVariants {\n Answer(AnswerResponseVariant),\n Fail(FailResponseVariant),\n}\n\n#[derive(Variant, serde::Deserialize)]\n#[variant(\n variant = \"Answer\",\n scenario = \"You have reviewed the blog post\",\n description = \"Suggestions for improving the blog post\"\n)]\npub struct AnswerResponseVariant {\n #[schema(\n description = \"Suggestions for improving the blog post\",\n example = \"[\\\"You wrote 'excellent' in two consecutive paragraphs in section 'Introduction'\\\"]\"\n )]\n pub suggestions: Vec,\n}\n\n#[derive(Variant, serde::Deserialize)]\n#[variant(\n variant = \"Fail\",\n scenario = \"For some reason you failed to generate suggestions\",\n description = \"Reason why you failed to generate suggestions\"\n)]\npub struct FailResponseVariant {\n #[schema(\n description = \"Reason why you failed to generate suggestions\",\n example = \"Content was invalid\"\n )]\n pub reason: String,\n}\n\n#[tokio::main]\nasync fn main() {\n let (lm, _) = get_lm();\n\n // In this example, we use the same LLM for the correction as for the main task.\n // This could be replaced by a smaller LM.\n let (corrector_lm, _) = get_lm();\n\n // We define an alignment strategy that uses the correction model.\n let alignment_strategy = AlignmentStrategyBuilder::new()\n .with_retries(2)\n .with_lm(&*corrector_lm)\n .try_build()\n .unwrap();\n\n // Mock blog post\n let prompt = \"\n This is a blog post about the importance of blogging.\n\n # Introduction\n\n Blogging is a crucial skill for any writer. It allows you to share your thoughts and ideas with others, and it can help you build a following and establish yourself as an expert in your field.\n \";\n\n let executor = StructuredExecutorBuilder::new()\n .with_lm(&*lm)\n .with_preamble(\"\n You are an experienced writer and blog post reviewer who helps users improve their blog posts.\n You will receive a blog post written in Markdown, and you will need to provide suggestions for improving it.\n Provide *specific* suggestions for improving the blog post, these can as nitpicky as you want.\n Consider things such as grammar, spelling, clarity, and conciseness.\n Even things like mentioning the same phrase too much in one paragraph, etc.\n The tone should be personal, friendly and professional at the same time.\n\n Be very specific and refer to specific sentences, paragraph and sections of the blog post.\n \")\n .with_options(Box::new(variants!(ResponseVariants)))\n .with_alignment(alignment_strategy)\n .try_build()\n .unwrap();\n let response = executor.execute(prompt).await.expect(\"Execution failed\");\n\n match response.content {\n ResponseVariants::Answer(answer) => {\n assert!(!answer.suggestions.is_empty());\n\n println!(\"Suggestions for improving the blog post:\");\n for suggestion in answer.suggestions {\n println!(\"- {}\", suggestion);\n }\n }\n ResponseVariants::Fail(fail) => {\n eprintln!(\"Model failed to generate a response: {}\", fail.reason);\n std::process::exit(1);\n }\n }\n}\n" }, { "path": "core/examples/structured_data_generation_capital.rs", "content": "//! This example demonstrates how to use the `Executor` to generate a structured response from the LLM.\n//! Run like so: `cargo run --example structured_data_generation_capital -- France`\n\n#![allow(dead_code)]\n\nuse orch::execution::*;\nuse orch::response::*;\n\nmod example_utils;\nuse example_utils::get_lm;\n\n#[derive(Variants, serde::Deserialize)]\npub enum ResponseVariants {\n Answer(AnswerResponseVariant),\n Fail(FailResponseVariant),\n}\n\n#[derive(Variant, serde::Deserialize)]\n#[variant(\n variant = \"Answer\",\n scenario = \"You know the capital city of the country\",\n description = \"Capital city of the country\"\n)]\npub struct AnswerResponseVariant {\n #[schema(\n description = \"Capital city of the received country\",\n example = \"London\"\n )]\n pub capital: String,\n}\n\n#[derive(Variant, serde::Deserialize)]\n#[variant(\n variant = \"Fail\",\n scenario = \"You don't know the capital city of the country\",\n description = \"Reason why the capital city is not known\"\n)]\npub struct FailResponseVariant {\n #[schema(\n description = \"Reason why the capital city is not known\",\n example = \"Country 'foobar' does not exist\"\n )]\n pub reason: String,\n}\n\n#[tokio::main]\nasync fn main() {\n let (lm, _) = get_lm();\n\n let country = \"France\";\n\n let executor = StructuredExecutorBuilder::new()\n .with_lm(&*lm)\n .with_preamble(\"\n You are a geography expert who helps users understand the capital city of countries around the world.\n You will receive a country name, and you will need to provide the capital city of that country.\n \")\n .with_options(Box::new(variants!(ResponseVariants)))\n .try_build()\n .unwrap();\n let response = executor.execute(country).await.expect(\"Execution failed\");\n\n match response.content {\n ResponseVariants::Answer(answer) => {\n println!(\"Capital city of {}: {}\", country, answer.capital);\n assert_eq!(answer.capital, \"Paris\");\n }\n ResponseVariants::Fail(fail) => {\n println!(\"Model failed to generate a response: {}\", fail.reason);\n }\n }\n}\n" }, { "path": "core/examples/text_generation.rs", "content": "//! This example demonstrates how to use the `Executor` to generate a response from the LLM.\n//! Run like so: `cargo run --example text_generation`\n\nuse orch::execution::*;\n\nmod example_utils;\nuse example_utils::get_lm;\n\n#[tokio::main]\nasync fn main() {\n let (lm, _) = get_lm();\n\n let prompt = \"What is 2+2?\";\n\n println!(\"Prompt: {prompt}\");\n println!(\"---\");\n\n let executor = TextExecutorBuilder::new()\n .with_lm(&*lm)\n .try_build()\n .unwrap();\n let response = executor.execute(prompt).await.expect(\"Execution failed\");\n\n println!(\"Response:\");\n println!(\"{}\", response.content);\n\n assert!(response.content.contains('4'));\n}\n" }, { "path": "core/examples/text_generation_stream.rs", "content": "//! This example demonstrates how to use the `Executor` to generate a streaming response from the LLM.\n//! Run like so: `cargo run --example text_generation_stream`\n\nuse orch::{execution::*, lm::LanguageModelProvider};\nuse tokio_stream::StreamExt;\n\nmod example_utils;\nuse example_utils::get_lm;\n\n#[tokio::main]\nasync fn main() {\n let (lm, provider) = get_lm();\n\n if provider == LanguageModelProvider::Anthropic {\n println!(\"Streaming is not currently supported for Anthropic. Skipping example.\");\n return;\n }\n\n let prompt = \"What is 2+2?\";\n\n println!(\"Prompt: {prompt}\");\n println!(\"---\");\n\n let executor = TextExecutorBuilder::new()\n .with_lm(&*lm)\n .try_build()\n .unwrap();\n let mut response = executor\n .execute_stream(prompt)\n .await\n .expect(\"Execution failed\");\n\n let mut response_text = String::new();\n println!(\"Response:\");\n while let Some(chunk) = response.stream.next().await {\n match chunk {\n Ok(chunk) => {\n print!(\"{chunk}\");\n response_text.push_str(&chunk);\n }\n Err(e) => {\n println!(\"Error: {e}\");\n break;\n }\n }\n }\n println!();\n\n assert!(!response_text.is_empty());\n}\n" }, { "path": "core/examples/variants_derive.rs", "content": "//! This example demonstrates how to use the `Variants` derive macro to generate a structured response from the LLM.\n//!\n//! Run like so: `cargo run --example variants_derive`\n\nuse orch::response::*;\n\n#[derive(Variants, serde::Deserialize)]\npub enum ResponseOptions {\n Answer(AnswerResponseOption),\n Fail(FailResponseOption),\n}\n\n#[derive(Variant, serde::Deserialize)]\n#[variant(\n variant = \"Answer\",\n scenario = \"You know the capital city of the country\",\n description = \"Capital city of the country\"\n)]\npub struct AnswerResponseOption {\n #[schema(\n description = \"Capital city of the received country\",\n example = \"London\"\n )]\n pub capital: String,\n #[schema(\n description = \"Country of the received capital city\",\n example = \"United Kingdom\"\n )]\n pub country: String,\n}\n\n#[derive(Variant, serde::Deserialize)]\n#[variant(\n variant = \"Fail\",\n scenario = \"You don't know the capital city of the country\",\n description = \"Reason why the capital city is not known\"\n)]\npub struct FailResponseOption {\n #[schema(\n description = \"Reason why the capital city is not known\",\n example = \"Country 'foobar' does not exist\"\n )]\n pub reason: String,\n}\n\nfn main() {\n let response = r#\"\n {\n \"response_type\": \"Answer\",\n \"capital\": \"London\",\n \"country\": \"United Kingdom\"\n }\n \"#;\n let parsed_response = variants!(ResponseOptions).parse(response).unwrap();\n match parsed_response {\n ResponseOptions::Answer(answer_response) => {\n println!(\"{}\", answer_response.capital);\n }\n ResponseOptions::Fail(fail_response) => {\n println!(\"{}\", fail_response.reason);\n }\n }\n}\n" }, { "path": "core/src/alignment/mod.rs", "content": "//! A module containing all logic related to alignment.\n//! Alignment, in this context, means \"aligning\" the model's output with the desired output.\n//! This takes the form of a so-called `[AlignmentStrategy]`, which is a trait that defines how to align the model's output.\n//!\n//! This concept has similarities to to traditional \"resilience\" techniques and libraries, such as .NET's [Polly](https://github.com/App-vNext/Polly),\n//! which I personally like a lot.\n\nmod strategy;\nmod strategy_builder;\n\npub use strategy::*;\npub use strategy_builder::*;\n" }, { "path": "core/src/alignment/strategy.rs", "content": "use async_recursion::async_recursion;\nuse orch_response_derive::{variants, Variant, Variants};\nuse thiserror::Error;\n\nuse crate::{\n execution::{ExecutorError, StructuredExecutor, StructuredExecutorBuilder},\n lm::{LanguageModel, LanguageModelError, TextCompleteOptions},\n};\n\n#[derive(Debug, Error)]\npub enum AlignmentError {\n #[error(\"Execution failed: {0}\")]\n ExecutionFailed(String),\n\n #[error(\"Language model error: {0}\")]\n LanguageModelError(#[from] LanguageModelError),\n\n #[error(\"Internal error: {0}\")]\n InternalError(String),\n\n #[error(\"Max retries exceeded ({0} retries)\")]\n MaxRetriesExceeded(usize),\n}\n\npub struct AlignmentStrategy<'a> {\n pub(crate) lm: &'a dyn LanguageModel,\n pub(crate) retries: usize,\n}\n\n#[derive(Variants, Clone, serde::Deserialize)]\npub enum AlignmentResponse {\n ResponseCorrection(ResponseCorrectionResponseVariant),\n SchemaCorrection(SchemaCorrectionResponseVariant),\n NoCorrection(NoCorrectionResponseVariant),\n Fail(FailResponseVariant),\n}\n\n#[derive(Variant, Clone, serde::Deserialize)]\n#[variant(\n variant = \"ResponseCorrection\",\n scenario = \"The response format is correct, but the response content itself is incorrect\",\n description = \"A correction and a reason why it is needed\"\n)]\npub struct ResponseCorrectionResponseVariant {\n #[schema(\n description = \"Correction of the phrase\",\n example = \"{ \\\"capital\\\": \\\"Paris\\\" }\"\n )]\n pub correction: String,\n\n #[schema(\n description = \"Short reason why a correction is needed\",\n example = \"The capital of France is not London as the original model returned, but Paris\"\n )]\n pub reason: String,\n}\n\n#[derive(Variant, Clone, serde::Deserialize)]\n#[variant(\n variant = \"SchemaCorrection\",\n scenario = \"The schema of the response is incorrect\",\n description = \"Explanation of why the schema is incorrect\"\n)]\npub struct SchemaCorrectionResponseVariant {\n #[schema(\n description = \"Correction of the schema, in natural language\",\n example = \"\\\"'capital' should be a string, not a number'\\\" or \\\"The 'capital' field has a typo and starts with an uppercase letter\\\"\"\n )]\n pub correction: String,\n\n #[schema(\n description = \"Short reason why a correction is needed\",\n example = \"The 'capital' field is a number, not a string\"\n )]\n pub reason: String,\n}\n\n#[derive(Variant, Clone, serde::Deserialize)]\n#[variant(\n variant = \"NoCorrection\",\n scenario = \"No correction needed, the original response satisfies the expected output\",\n description = \"Short reason why a correction is not needed\"\n)]\npub struct NoCorrectionResponseVariant {\n #[schema(\n description = \"Short reason why a correction is not needed\",\n example = \"The user asked for the capital city of France, and the answer is indeed Paris\"\n )]\n pub reason: String,\n}\n\n#[derive(Variant, Clone, serde::Deserialize)]\n#[variant(\n variant = \"Fail\",\n scenario = \"You don't know how to verify whether the answer is correct or not. You should only go for this response in extreme cases\",\n description = \"Reason why you failed to determine whether the answer is correct or not\"\n)]\npub struct FailResponseVariant {\n #[schema(\n description = \"Reason why you failed to determine whether the answer is correct or not\",\n example = \"The question is extremely vague and the model returned something completely unrelated\"\n )]\n pub reason: String,\n}\n\nimpl<'a> AlignmentStrategy<'a> {\n const PREAMBLE: &'static str = \"\n Your purpose is to receive a response from a language model and make sure (and correct otherwise) whether the response is expected or not. \n Being \\\"expected\\\" means that the response is correct and matches the expected output.\n\n You should *not* return the response in the schema of the original message, but instead of the schema that you are requested to provide\n (the one with the response types 'ResponseCorrection', 'SchemaCorrection' and 'NoCorrection').\n \";\n\n /// Aligns the response of the language model.\n /// Tries at least once, and continues according to the [`AlignmentStrategy`]\n /// (e.g., number of retries).\n pub async fn align(\n &self,\n base_lm: &'a dyn LanguageModel,\n original_preamble: &str,\n original_prompt: &str,\n original_response: &str,\n ) -> Result {\n let mut iterated_response = original_response.to_owned();\n let mut retry_count = 0;\n let mut prev_alignment_response = None;\n\n loop {\n let response = self\n .request_correction(\n original_preamble,\n original_prompt,\n &iterated_response,\n &prev_alignment_response,\n )\n .await?;\n\n let Some(response) = response else {\n // The response may be `None` if the correction deemed that the previous response should be used.\n continue;\n };\n\n match &response {\n AlignmentResponse::NoCorrection(_) => {\n // Found no correction, can return the original response.\n return Ok(iterated_response.to_owned());\n }\n response => {\n retry_count += 1;\n\n if retry_count >= self.retries {\n return Err(AlignmentError::MaxRetriesExceeded(retry_count));\n }\n\n if let AlignmentResponse::Fail(_) = response {\n // Failed - simply try again.\n continue;\n }\n\n let correction = match response {\n AlignmentResponse::ResponseCorrection(response_correction) => {\n response_correction.correction.clone()\n }\n AlignmentResponse::SchemaCorrection(schema_correction) => {\n schema_correction.correction.clone()\n }\n _ => unreachable!(),\n };\n\n let correction_prompt = format!(\"\n {original_preamble}\n\n NOTE:\n You have previously answered this with the following response and was incorrect. Here is the response and the correction, please make sure not to repeat the same mistake:\n ORIGINAL RESPONSE: {original_response}\n CORRECTION: {correction}\n \");\n let new_base_model_response = base_lm\n .text_complete(\n original_prompt,\n &correction_prompt,\n TextCompleteOptions::default(),\n )\n .await\n .map_err(AlignmentError::LanguageModelError)?;\n\n prev_alignment_response = Some(response.clone());\n iterated_response = new_base_model_response.text;\n }\n }\n }\n }\n\n #[async_recursion]\n async fn request_correction(\n &self,\n original_preamble: &str,\n original_prompt: &str,\n original_response: &str,\n prev_alignment_response: &Option,\n ) -> Result, AlignmentError> {\n let mut preamble = format!(\n \"\n {base_preamble}\n\n The model received the original instructions:\n {original_preamble}\n\n And the original prompt:\n {original_prompt}\n\n And the original response:\n {original_response}\n\n REMEMBER: Return a response in the schema you are requested (the one with the response types 'ResponseCorrection', 'SchemaCorrection' and 'NoCorrection').\n \",\n base_preamble = Self::PREAMBLE,\n );\n // If `alignment_response` is `None`, then this was the first attempt and no additional preamble is needed.\n if let Some(prev_alignment_response) = prev_alignment_response {\n // TODO: Add context of more tries?\n preamble.push_str(&format!(\n \"\n IMPORTANT CONTEXT:\n Before receiving the previous correction, the model has already responded with the following:\n\n {}\n\n And received the following corrections:\n \",\n original_response\n ));\n\n match prev_alignment_response {\n AlignmentResponse::ResponseCorrection(response_correction) => {\n preamble.push_str(&format!(\n \"CORRECTION: The response content was incorrect, this is the correction: {}\",\n response_correction.correction,\n ));\n }\n AlignmentResponse::SchemaCorrection(schema_correction) => {\n preamble.push_str(&format!(\n \"CORRECTION: The response schema was incorrect for the following reason: {}\n This is the correction: {}\n \",\n schema_correction.correction, schema_correction.reason\n ));\n }\n _ => {\n // No error (this is unexpected) - return the original response.\n return Err(AlignmentError::InternalError(\n \"Requested correction with no relevant correction response\".to_owned(),\n ));\n }\n }\n }\n\n let executor: StructuredExecutor = StructuredExecutorBuilder::new()\n .with_lm(self.lm)\n .with_preamble(&preamble)\n .with_options(Box::new(variants!(AlignmentResponse)))\n .try_build()\n .unwrap();\n let response = {\n let correction_response = executor.execute(original_prompt).await;\n\n match correction_response {\n Ok(response) => Some(response.content),\n Err(ExecutorError::Parsing(_)) => {\n // The model failed to parse the response, so we return the original response.\n return Ok(None);\n }\n Err(e) => return Err(AlignmentError::ExecutionFailed(e.to_string())),\n }\n };\n\n Ok(response)\n }\n}\n" }, { "path": "core/src/alignment/strategy_builder.rs", "content": "use thiserror::Error;\n\nuse crate::lm::LanguageModel;\n\nuse super::strategy::AlignmentStrategy;\n\n/// The default number of retries for the alignment strategy, if not overriden.\npub const DEFAULT_RETRIES: usize = 2;\n\n#[derive(Debug, Error)]\npub enum AlignmentStrategyBuilderError {\n #[error(\"{0} is not set\")]\n ConfigurationNotSet(String),\n}\n\n#[derive(Default)]\npub struct AlignmentStrategyBuilder<'a> {\n lm: Option<&'a dyn LanguageModel>,\n retries: Option,\n}\n\nimpl<'a> AlignmentStrategyBuilder<'a> {\n /// Creates a new `AlignmentStrategyBuilder` instance.\n pub fn new() -> Self {\n Self {\n lm: None,\n retries: Some(DEFAULT_RETRIES),\n }\n }\n\n /// Sets the language model to use for the alignment strategy.\n pub fn with_lm(mut self, lm: &'a dyn LanguageModel) -> Self {\n self.lm = Some(lm);\n self\n }\n\n /// Sets the number of retries for the alignment strategy.\n pub fn with_retries(mut self, retries: usize) -> Self {\n self.retries = Some(retries);\n self\n }\n\n /// Builds the alignment strategy.\n /// May fail with a [`AlignmentStrategyBuilderErrro`] if some required configurations are not set.\n pub fn try_build(self) -> Result, AlignmentStrategyBuilderError> {\n let Some(lm) = self.lm else {\n return Err(AlignmentStrategyBuilderError::ConfigurationNotSet(\n \"Language model\".to_string(),\n ));\n };\n let Some(retries) = self.retries else {\n return Err(AlignmentStrategyBuilderError::ConfigurationNotSet(\n \"Retries\".to_string(),\n ));\n };\n Ok(AlignmentStrategy { lm, retries })\n }\n}\n" }, { "path": "core/src/execution/builder.rs", "content": "use thiserror::Error;\n\n#[derive(Debug, Error)]\npub enum ExecutorBuilderError {\n #[error(\"Internal error: {0}\")]\n InternalError(String),\n #[error(\"{0} is not set\")]\n ConfigurationNotSet(String),\n}\n" }, { "path": "core/src/execution/executor.rs", "content": "use std::pin::Pin;\n\nuse thiserror::Error;\nuse tokio_stream::Stream;\n\nuse crate::{\n alignment::AlignmentError,\n lm::{LanguageModel, LanguageModelError, OllamaError, TextCompleteOptions},\n};\n\n#[derive(Debug, Error)]\npub enum ExecutorError {\n #[error(\"{0}\")]\n General(LanguageModelError),\n\n #[error(\"{0}\")]\n LanguageModelError(LanguageModelError),\n\n #[error(\"Error when calling Ollama API: {0}\")]\n OllamaApi(String),\n\n #[error(\"Parsing LM response failed: {0}\")]\n Parsing(String),\n\n #[error(\"Alignment error: {0}\")]\n Alignment(AlignmentError),\n}\n\nimpl From for ExecutorError {\n fn from(val: LanguageModelError) -> Self {\n match val {\n LanguageModelError::Ollama(OllamaError::Api(e)) => ExecutorError::OllamaApi(e),\n e => ExecutorError::LanguageModelError(e),\n }\n }\n}\n\npub(crate) trait Executor<'a> {\n /// Generates a text completion from the LLM (non-streaming).\n async fn text_complete(\n &self,\n prompt: &str,\n ) -> Result, ExecutorError> {\n text_complete(self.lm(), prompt, &self.system_prompt()).await\n }\n\n /// System prompt (instructions) for the model.\n fn system_prompt(&self) -> String;\n\n fn lm(&self) -> &'a dyn LanguageModel;\n}\n\n// TODO: Support context for completions (e.g., IDs of past conversations in Ollama).\npub struct ExecutorContext;\n\npub struct ExecutorTextCompleteResponse {\n pub content: T,\n pub context: ExecutorContext,\n}\n\npub struct ExecutorTextCompleteStreamResponse {\n pub stream: Pin> + Send>>,\n pub context: ExecutorContext,\n}\n\npub async fn text_complete<'a>(\n lm: &'a dyn LanguageModel,\n prompt: &str,\n system_prompt: &str,\n) -> Result, ExecutorError> {\n let options = TextCompleteOptions {\n ..Default::default()\n };\n let response = lm\n .text_complete(prompt, system_prompt, options)\n .await\n .map_err(ExecutorError::from)?;\n Ok(ExecutorTextCompleteResponse {\n content: response.text,\n context: ExecutorContext {},\n })\n}\n\npub(crate) async fn generate_embedding<'a>(\n lm: &'a dyn LanguageModel,\n prompt: &str,\n) -> Result, ExecutorError> {\n let response = lm\n .generate_embedding(prompt)\n .await\n .map_err(ExecutorError::from)?;\n Ok(response)\n}\n" }, { "path": "core/src/execution/mod.rs", "content": "//! A module containing all logic related to LLM execution.\n//! An [`Executor`] is the terminology for a component which executes an LLM,\n//! and aligns it appropriately (e.g., error correction).\n//!\n//! It is not to be confused with an [`Orchestrator`] which manages the execution of an LLM\n//! or multiple LLMs towards a task.\n\nmod builder;\nmod executor;\nmod response;\nmod structured_executor;\nmod text_executor;\n\npub use builder::*;\npub use executor::*;\npub use response::*;\npub use structured_executor::*;\npub use text_executor::*;\n" }, { "path": "core/src/execution/response.rs", "content": "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum ResponseFormat {\n Text,\n Json,\n}\n\nimpl Default for ResponseFormat {\n fn default() -> Self {\n Self::Text\n }\n}\n" }, { "path": "core/src/execution/structured_executor.rs", "content": "use std::cell::OnceCell;\n\nuse orch_response::{OrchResponseVariants, ResponseSchemaField};\n\nuse crate::{alignment::AlignmentStrategy, lm::LanguageModel};\n\nuse super::{\n generate_embedding, Executor, ExecutorBuilderError, ExecutorContext, ExecutorError,\n ExecutorTextCompleteResponse, DEFAULT_PREAMBLE,\n};\n\npub struct StructuredExecutor<'a, T> {\n pub(crate) lm: &'a dyn LanguageModel,\n pub(crate) preamble: Option<&'a str>,\n pub(crate) variants: Box>,\n pub(crate) alignment_strategy: Option>,\n}\n\nimpl<'a, T> Executor<'a> for StructuredExecutor<'a, T> {\n fn lm(&self) -> &'a dyn LanguageModel {\n self.lm\n }\n\n fn system_prompt(&self) -> String {\n let cell = OnceCell::new();\n\n cell.get_or_init(|| {\n let response_options = self.variants.variants();\n let all_types = response_options\n .iter()\n .map(|option| option.type_name.clone())\n .collect::>();\n\n let response_options_text = response_options\n .iter()\n .map(|option| {\n let mut schema_text = String::new();\n let mut schema_example = \"{\".to_string();\n let type_field = ResponseSchemaField {\n // NOTE: This is assumed by [`orch_response_derive`] to be the discriminator field.\n name: \"response_type\".to_string(),\n description: format!(\n \"The type of the response (\\\"{}\\\" in this case)\",\n option.type_name\n )\n .to_string(),\n typ: \"string\".to_string(),\n example: all_types.first().unwrap().to_string(),\n };\n\n for (i, field) in option\n .schema\n .iter()\n .chain(std::iter::once(&type_field))\n .enumerate()\n {\n schema_text.push_str(&format!(\n \" - `{}` of type {} (description: {})\\n\\n\",\n field.name, field.typ, field.description\n ));\n schema_example\n .push_str(&format!(\"\\\"{}\\\": \\\"{}\\\"\", field.name, field.example));\n\n if i < option.schema.len() - 1 {\n schema_example.push(',');\n }\n }\n schema_example.push('}');\n\n format!(\n \"SCENARIO: {}\\nDESCRIPTION: {}\\nSCHEMA:\\n{}\\nEXAMPLE RESPONSE: {}\\n\\n\\n\",\n option.scenario, option.description, schema_text, schema_example\n )\n })\n .collect::>()\n .join(\"\\n\");\n\n // Add an optional extra preamble supplied by the user.\n let preamble = self.preamble.map(|pa| format!(\"Additional information: {}\", pa)).unwrap_or(\"\".to_owned());\n\n let system_prompt = format!(\n \"\n\t\t\t\t\t\t\t\tYou will receive a prompt from a user, and will need to response with a JSON object that represents the response.\n\t\t\t\t\t\t\t\tResponse *only* with the JSON object, and nothing else. No additional preamble or explanations. Only work with the responses you can reply with.\n\t\t\t\t\t\t\t\t{preamble}\n\n\t\t\t\t\t\t\t\tYou have {choices_len} choices to respond, in a JSON format:\n {response_options_text}\n \",\n preamble = preamble,\n choices_len = response_options.len(),\n response_options_text = response_options_text\n )\n .trim()\n .to_string();\n\n system_prompt\n })\n .clone()\n }\n}\n\n/// Trait for LLM execution.\n/// This should be implemented for each LLM text generation use-case, where the system prompt\n/// changes according to the trait implementations.\nimpl<'a, T> StructuredExecutor<'a, T> {\n /// Generates a structured response from the LLM (non-streaming).\n ///\n /// # Arguments\n /// * `prompt` - The prompt to generate a response for.\n /// * `system_prompt` - The system prompt to use for the generation.\n ///\n /// # Returns\n /// A [Result] containing the response from the LLM or an error if there was a problem.\n pub async fn execute(\n &'a self,\n prompt: &'a str,\n ) -> Result, ExecutorError> {\n let mut model_response = self.text_complete(prompt).await?.content;\n if let Some(alignment_strategy) = &self.alignment_strategy {\n model_response = alignment_strategy\n .align(\n self.lm,\n self.preamble.unwrap_or(DEFAULT_PREAMBLE),\n prompt,\n &model_response,\n )\n .await\n .map_err(ExecutorError::Alignment)?;\n }\n let result = self\n .variants\n .parse(&model_response)\n .map_err(|e| ExecutorError::Parsing(format!(\"{e}\\nResponse: {:?}\", model_response)))?;\n // TODO: Add error correction and handling.\n Ok(ExecutorTextCompleteResponse {\n content: result,\n context: ExecutorContext {},\n })\n }\n\n /// Generates an embedding from the LLM.\n ///\n /// # Arguments\n /// * `prompt` - The item to generate an embedding for.\n ///\n /// # Returns\n ///\n /// A [Result] containing the embedding or an error if there was a problem.\n pub async fn generate_embedding(&'a self, prompt: &'a str) -> Result, ExecutorError> {\n generate_embedding(self.lm, prompt).await\n }\n}\n\n#[derive(Default)]\npub struct StructuredExecutorBuilder<'a, T> {\n lm: Option<&'a dyn LanguageModel>,\n preamble: Option<&'a str>,\n variants: Option>>,\n alignment_strategy: Option>,\n}\n\nimpl<'a, T> StructuredExecutorBuilder<'a, T> {\n pub fn new() -> Self {\n Self {\n lm: None,\n preamble: None,\n variants: None,\n alignment_strategy: None,\n }\n }\n\n pub fn with_lm(mut self, lm: &'a dyn LanguageModel) -> Self {\n self.lm = Some(lm);\n self\n }\n\n pub fn with_options(mut self, options: Box>) -> Self {\n self.variants = Some(options);\n self\n }\n\n pub fn with_preamble(mut self, preamble: &'a str) -> Self {\n self.preamble = Some(preamble);\n self\n }\n\n pub fn with_alignment(mut self, strategy: AlignmentStrategy<'a>) -> Self {\n self.alignment_strategy = Some(strategy);\n self\n }\n\n pub fn try_build(self) -> Result, ExecutorBuilderError> {\n let Some(lm) = self.lm else {\n return Err(ExecutorBuilderError::ConfigurationNotSet(\n \"Language model\".to_string(),\n ));\n };\n let Some(response_options) = self.variants else {\n return Err(ExecutorBuilderError::InternalError(\n \"Response variants are not set\".to_string(),\n ));\n };\n Ok(StructuredExecutor {\n lm,\n preamble: self.preamble,\n variants: response_options,\n alignment_strategy: self.alignment_strategy,\n })\n }\n}\n" }, { "path": "core/src/execution/text_executor.rs", "content": "use crate::lm::{LanguageModel, TextCompleteStreamOptions};\n\nuse super::{\n generate_embedding, Executor, ExecutorBuilderError, ExecutorContext, ExecutorError,\n ExecutorTextCompleteResponse, ExecutorTextCompleteStreamResponse,\n};\n\npub const DEFAULT_PREAMBLE: &str = \"You are a helpful assistant\";\n\npub struct TextExecutor<'a> {\n pub(crate) lm: &'a dyn LanguageModel,\n pub(crate) preamble: Option<&'a str>,\n}\n\nimpl<'a> Executor<'a> for TextExecutor<'a> {\n fn lm(&self) -> &'a dyn LanguageModel {\n self.lm\n }\n\n fn system_prompt(&self) -> String {\n self.preamble.unwrap_or(DEFAULT_PREAMBLE).to_owned()\n }\n}\n\n/// Trait for LLM execution.\n/// This should be implemented for each LLM text generation use-case, where the system prompt\n/// changes according to the trait implementations.\nimpl<'a> TextExecutor<'a> {\n /// Generates a streaming response from the LLM.\n ///\n /// # Arguments\n /// * `prompt` - The prompt to generate a response for.\n /// * `system_prompt` - The system prompt to use for the generation.\n ///\n /// # Returns\n /// A [Result] containing the response from the LLM or an error if there was a problem.\n pub async fn execute_stream(\n &'a self,\n prompt: &'a str,\n ) -> Result {\n let options = TextCompleteStreamOptions {\n ..Default::default()\n };\n let system_prompt = self.system_prompt();\n let response = self\n .lm\n .text_complete_stream(prompt, &system_prompt, options)\n .await\n .map_err(ExecutorError::General)?;\n Ok(ExecutorTextCompleteStreamResponse {\n stream: response.stream,\n context: ExecutorContext {},\n })\n }\n\n /// Generates a response from the LLM (non-streaming).\n ///\n /// # Arguments\n /// * `prompt` - The prompt to generate a response for.\n /// * `system_prompt` - The system prompt to use for the generation.\n ///\n /// # Returns\n /// A [Result] containing the response from the LLM or an error if there was a problem.\n pub async fn execute(\n &'a self,\n prompt: &'a str,\n ) -> Result, ExecutorError> {\n self.text_complete(prompt).await\n }\n\n /// Generates an embedding from the LLM.\n ///\n /// # Arguments\n /// * `prompt` - The item to generate an embedding for.\n ///\n /// # Returns\n ///\n /// A [Result] containing the embedding or an error if there was a problem.\n pub async fn generate_embedding(&'a self, prompt: &'a str) -> Result, ExecutorError> {\n generate_embedding(self.lm, prompt).await\n }\n}\n\n#[derive(Default)]\npub struct TextExecutorBuilder<'a> {\n lm: Option<&'a dyn LanguageModel>,\n preamble: Option<&'a str>,\n}\n\nimpl<'a> TextExecutorBuilder<'a> {\n pub fn new() -> Self {\n Self {\n lm: None,\n preamble: None,\n }\n }\n\n pub fn with_lm(mut self, lm: &'a dyn LanguageModel) -> Self {\n self.lm = Some(lm);\n self\n }\n\n pub fn with_preamble(mut self, preamble: &'a str) -> Self {\n self.preamble = Some(preamble);\n self\n }\n\n pub fn try_build(self) -> Result, ExecutorBuilderError> {\n let Some(lm) = self.lm else {\n return Err(ExecutorBuilderError::ConfigurationNotSet(\n \"Language model\".to_string(),\n ));\n };\n Ok(TextExecutor {\n lm,\n preamble: self.preamble,\n })\n }\n}\n" }, { "path": "core/src/lib.rs", "content": "#![doc = include_str!(concat!(env!(\"CARGO_MANIFEST_DIR\"), \"/README.md\"))]\n\npub mod alignment;\npub mod execution;\npub mod lm;\nmod net;\npub mod response;\n" }, { "path": "core/src/lm/builder.rs", "content": "use thiserror::Error;\n\nuse super::LanguageModel;\n\n#[derive(Debug, Error)]\npub enum LanguageModelBuilderError {\n #[error(\"{0} is not set\")]\n ConfigurationNotSet(String),\n}\n\npub trait LanguageModelBuilder {\n fn new() -> Self;\n fn try_build(self) -> Result;\n}\n" }, { "path": "core/src/lm/error.rs", "content": "use thiserror::Error;\n\nuse super::{AnthropicError, LanguageModelProvider, OllamaError, OpenAiError};\n\n#[derive(Debug, Error)]\npub enum LanguageModelProviderError {\n #[error(\"Invalid LLM provider: {0}\")]\n InvalidValue(String),\n}\n\nimpl std::fmt::Display for LanguageModelProvider {\n fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {\n match self {\n LanguageModelProvider::Ollama => write!(f, \"ollama\"),\n LanguageModelProvider::OpenAi => write!(f, \"openai\"),\n LanguageModelProvider::Anthropic => write!(f, \"anthropic\"),\n }\n }\n}\n\nimpl Default for LanguageModelProvider {\n fn default() -> Self {\n Self::Ollama\n }\n}\n\n#[derive(Debug, Error)]\npub enum LanguageModelError {\n #[error(\"Text generation error: {0}\")]\n TextGeneration(String),\n\n #[error(\"Feature unsupported: {0}\")]\n UnsupportedFeature(String),\n\n #[error(\"Embedding generation error: {0}\")]\n EmbeddingGeneration(String),\n\n #[error(\"Configuration error: {0}\")]\n Configuration(String),\n\n #[error(\"Ollama error: {0}\")]\n Ollama(#[from] OllamaError),\n\n #[error(\"OpenAI error: {0}\")]\n OpenAi(#[from] OpenAiError),\n\n #[error(\"Anthropic error: {0}\")]\n Anthropic(#[from] AnthropicError),\n}\n" }, { "path": "core/src/lm/lm_provider/anthropic/builder.rs", "content": "use thiserror::Error;\n\nuse crate::lm::{LanguageModelBuilder, LanguageModelBuilderError};\n\nuse super::client::config::{DEFAULT_API_ENDPOINT, DEFAULT_MODEL};\nuse super::Anthropic;\n\n#[derive(Debug, Error)]\npub enum AnthropicBuilderError {\n #[error(\"Configuration error: {0} is not set\")]\n ConfigurationNotSet(String),\n}\n\n/// Builds an [`Anthropic`] instance.\npub struct AnthropicBuilder {\n /// API key for the Anthropic API. Required.\n api_key: Option,\n /// Base URL for the Anthropic API. Defaults to [`DEFAULT_BASE_URL`].\n api_endpoint: Option,\n /// Model to use for text completion. Defaults to [`DEFAULT_MODEL`].\n model: Option,\n}\n\nimpl AnthropicBuilder {\n /// Overrides the default base URL for the Anthropic API.\n /// Defaults to [`DEFAULT_API_ENDPOINT`].\n pub fn with_api_endpoint(mut self, base_url: String) -> Self {\n self.api_endpoint = Some(base_url);\n self\n }\n\n /// Sets the required API key for the Anthropic API.\n pub fn with_api_key(mut self, api_key: String) -> Self {\n self.api_key = Some(api_key);\n self\n }\n\n /// Overrides the default model to use for text completion.\n /// Defaults to [`DEFAULT_MODEL`].\n pub fn with_model(mut self, model: String) -> Self {\n self.model = Some(model);\n self\n }\n}\n\nimpl LanguageModelBuilder for AnthropicBuilder {\n fn new() -> Self {\n Self {\n api_key: None,\n api_endpoint: Some(DEFAULT_API_ENDPOINT.to_string()),\n model: Some(DEFAULT_MODEL.to_string()),\n }\n }\n\n /// Tries to build an [`Anthropic`] instance. May fail if the required configurations are not set.\n fn try_build(self) -> Result {\n let Some(api_endpoint) = self.api_endpoint else {\n return Err(LanguageModelBuilderError::ConfigurationNotSet(\n \"API endpoint\".to_string(),\n ));\n };\n let Some(api_key) = self.api_key else {\n return Err(LanguageModelBuilderError::ConfigurationNotSet(\n \"API key\".to_string(),\n ));\n };\n let Some(model) = self.model else {\n return Err(LanguageModelBuilderError::ConfigurationNotSet(\n \"Model\".to_string(),\n ));\n };\n Ok(Anthropic {\n api_key,\n api_endpoint,\n model: model.to_owned(),\n })\n }\n}\n" }, { "path": "core/src/lm/lm_provider/anthropic/client/anthropic_client.rs", "content": "#![allow(dead_code)]\n\nuse thiserror::Error;\n\nuse crate::lm::lm_provider::anthropic::client::models::AnthropicMessagesApiRequest;\n\nuse super::{\n config::DEFAULT_MAX_TOKENS,\n models::{\n AnthropicMessage, AnthropicMessagesApiMessage, AnthropicMessagesApiResponse,\n AnthropicMessagesApiResponseSuccess,\n },\n};\n\n#[derive(Debug, Error)]\npub(crate) enum AnthropicClientError {\n #[error(\"{0}\")]\n InternalError(String),\n\n #[error(\"Configuration '{0}' is not set\")]\n ConfigurationNotSet(String),\n\n #[error(\"Failed to serialize/deserialize: {0}\")]\n Marhsalling(String),\n\n #[error(\"Failed to send or receive request to/from Anthropic API: {0}\")]\n Api(String),\n}\n\n/// A client for interacting with the Anthropic API.\npub struct AnthropicClient {\n pub(crate) api_endpoint: String,\n pub(crate) api_key: String,\n}\n\nimpl AnthropicClient {\n /// Generates a response from the Anthropic API.\n ///\n /// # Arguments\n /// * `prompt` - The prompt to generate a response for. Use \"User:...\\n\\n\" for user messages and \"Assistant:...\\n\\n\" for assistant messages.\n /// * `system_prompt` - The system prompt to use for the generation.\n /// * `options` - The options for the generation (use [`AnthropicClientTextCompleteOptionsBuilder`] to build a new instance).\n ///\n /// # Returns\n /// A [Result] containing the response from the Anthropic API or an error if there was a problem.\n pub async fn text_complete(\n &self,\n messages: &[AnthropicMessage],\n system_prompt: &str,\n options: AnthropicClientTextCompleteOptions,\n ) -> Result {\n let messages_api_endpoint = format!(\"{}/v1/messages\", self.api_endpoint);\n\n let system_prompt = if system_prompt.is_empty() {\n None\n } else {\n Some(system_prompt.to_string())\n };\n\n let messages = messages\n .iter()\n .map(Self::construct_message)\n .collect::>();\n\n let req_body = AnthropicMessagesApiRequest {\n messages,\n system_prompt,\n model: options.model,\n max_tokens_to_sample: DEFAULT_MAX_TOKENS,\n stop_sequences: None,\n temperature: None,\n top_k: None,\n };\n\n let http_client = reqwest::Client::new();\n let req = http_client\n .post(messages_api_endpoint)\n // See Anthropic authentication documentation: https://docs.anthropic.com/en/api/getting-started#authentication\n .header(\"x-api-key\", &self.api_key)\n .header(\"anthropic-version\", \"2023-06-01\")\n .header(reqwest::header::CONTENT_TYPE, \"application/json\")\n .body(\n serde_json::to_string(&req_body)\n .map_err(|e| AnthropicClientError::Marhsalling(e.to_string()))?,\n )\n .build()\n .map_err(|e| AnthropicClientError::InternalError(e.to_string()))?;\n\n let response = http_client\n .execute(req)\n .await\n .map_err(|e| AnthropicClientError::Api(e.to_string()))?;\n\n let response_body_json = response\n .text()\n .await\n .map_err(|e| AnthropicClientError::Api(e.to_string()))?\n .to_string();\n\n let deserialized_response: AnthropicMessagesApiResponse =\n serde_json::from_str(&response_body_json).map_err(|e| {\n AnthropicClientError::Marhsalling(format!(\n \"Failed to parse response: {e} (response: {response_body_json})\"\n ))\n })?;\n\n Ok(match deserialized_response {\n AnthropicMessagesApiResponse::Success(success_response) => success_response,\n AnthropicMessagesApiResponse::Error(error_response) => {\n let error_message = error_response.error.message;\n return Err(AnthropicClientError::Api(error_message));\n }\n })\n }\n\n fn construct_message(msg: &AnthropicMessage) -> AnthropicMessagesApiMessage {\n match msg {\n AnthropicMessage::User(content) => AnthropicMessagesApiMessage {\n role: \"user\".to_string(),\n content: content.to_string(),\n },\n AnthropicMessage::Assistant(content) => AnthropicMessagesApiMessage {\n role: \"assistant\".to_string(),\n content: content.to_string(),\n },\n }\n }\n}\n\n/// Options for text completion.\n#[derive(Debug, Default)]\npub struct AnthropicClientTextCompleteOptions {\n /// See [`AnthropicCompleteApiRequest::model`].\n pub model: String,\n /// See [`AnthropicCompleteApiRequest::max_tokens`].\n pub max_tokens_to_sample: usize,\n /// See [`AnthropicCompleteApiRequest::stop_sequences`].\n pub stop_sequences: Option>,\n /// See [`AnthropicCompleteApiRequest::temperature`].\n pub temperature: Option,\n /// See [`AnthropicCompleteApiRequest::top_k`].\n pub top_k: Option,\n}\n\n/// Builds a new [`AnthropicClientTextCompleteOptions`] instance.\n#[derive(Debug, Default)]\npub struct AnthropicClientTextCompleteOptionsBuilder {\n model: Option,\n max_tokens: usize,\n stop_sequences: Option>,\n temperature: Option,\n top_k: Option,\n}\n\nimpl AnthropicClientTextCompleteOptionsBuilder {\n pub fn new() -> Self {\n Self {\n max_tokens: DEFAULT_MAX_TOKENS,\n ..Default::default()\n }\n }\n\n /// Sets the model (required).\n pub fn with_model(mut self, model: String) -> Self {\n self.model = Some(model);\n self\n }\n\n /// Sets the maximum number of tokens to generate before stopping.\n /// Defaults to [`DEFAULT_MAX_TOKENS`].\n pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {\n self.max_tokens = max_tokens;\n self\n }\n\n /// Sets the stop sequences.\n pub fn with_stop_sequences(mut self, stop_sequences: Vec) -> Self {\n self.stop_sequences = Some(stop_sequences);\n self\n }\n\n /// Sets the temperature.\n pub fn with_temperature(mut self, temperature: f32) -> Self {\n self.temperature = Some(temperature);\n self\n }\n\n /// Sets the top k.\n pub fn with_top_k(mut self, top_k: usize) -> Self {\n self.top_k = Some(top_k);\n self\n }\n\n /// Tries to build a [`AnthropicClientTextCompleteOptions`] instance. May fail if the required configurations are not set.\n pub fn try_build(self) -> Result {\n let Some(model) = self.model else {\n return Err(AnthropicClientError::ConfigurationNotSet(\n \"Model\".to_string(),\n ));\n };\n Ok(AnthropicClientTextCompleteOptions {\n model,\n max_tokens_to_sample: self.max_tokens,\n stop_sequences: self.stop_sequences,\n temperature: self.temperature,\n top_k: self.top_k,\n })\n }\n}\n" }, { "path": "core/src/lm/lm_provider/anthropic/client/builder.rs", "content": "use thiserror::Error;\n\nuse super::{anthropic_client::AnthropicClient, config};\n\n#[derive(Debug, Error)]\npub enum AnthropicBuilderError {\n #[error(\"Configuration error: {0} is not set\")]\n ConfigurationNotSet(String),\n}\n\n/// Builds an [`AnthropicClient`] instance.\npub struct AnthropicClientBuilder {\n api_endpoint: String,\n api_key: Option,\n}\n\nimpl AnthropicClientBuilder {\n pub fn new() -> Self {\n Self {\n api_endpoint: config::DEFAULT_API_ENDPOINT.to_string(),\n api_key: None,\n }\n }\n\n /// Sets an override for the Anthropic API endpoint. Defaults to [`config::DEFAULT_API_ENDPOINT`].\n pub fn with_api_endpoint(mut self, api_endpoint: String) -> Self {\n self.api_endpoint = api_endpoint;\n self\n }\n\n /// Sets the required Anthropic API key.\n pub fn with_api_key(mut self, api_key: String) -> Self {\n self.api_key = Some(api_key);\n self\n }\n\n pub fn try_build(self) -> Result {\n let Some(api_key) = self.api_key else {\n return Err(AnthropicBuilderError::ConfigurationNotSet(\n \"API key\".to_string(),\n ));\n };\n Ok(AnthropicClient {\n api_endpoint: self.api_endpoint,\n api_key,\n })\n }\n}\n" }, { "path": "core/src/lm/lm_provider/anthropic/client/config.rs", "content": "use crate::lm::anthropic_model;\n\n/// Default API endpoint for the Anthropic API.\npub const DEFAULT_API_ENDPOINT: &str = \"https://api.anthropic.com\";\n/// Default model to use for text completion.\npub const DEFAULT_MODEL: &str = anthropic_model::CLAUDE_3_5_SONNET;\n/// Default maximum number of tokens to generate before stopping.\npub const DEFAULT_MAX_TOKENS: usize = 2048;\n" }, { "path": "core/src/lm/lm_provider/anthropic/client/mod.rs", "content": "pub mod anthropic_client;\npub mod builder;\npub mod config;\npub mod models;\n" }, { "path": "core/src/lm/lm_provider/anthropic/client/models.rs", "content": "use serde::{Deserialize, Serialize};\n\n/// Request for generating a response from the Anthropic API.\n/// Referenced from the Anthropic API documentation [here](https://docs.anthropic.com/en/api/complete).\n#[derive(Debug, Serialize, Deserialize)]\npub struct AnthropicMessagesApiRequest {\n /// The model that will complete your prompt.\n /// See [`anthropic_model`] for a list of built-in model IDs for convenience.\n ///\n /// See [models](https://docs.anthropic.com/en/docs/about-claude/models) for a complete list of models supported by Anthropic.\n pub model: String,\n\n /// Messages to generate a completion for.\n ///\n /// See [Anthropic API documentation](https://docs.anthropic.com/en/api/messages) for more information.\n pub messages: Vec,\n\n /// Optional system prompt.\n #[serde(skip_serializing_if = \"Option::is_none\")]\n #[serde(rename = \"system\")]\n pub system_prompt: Option,\n\n /// The maximum number of tokens to generate before stopping.\n ///\n /// Note that the Anthropic models may stop before reaching this maximum. This parameter only specifies the absolute maximum number of tokens to generate.\n #[serde(rename = \"max_tokens\")]\n pub max_tokens_to_sample: usize,\n\n /// Sequences that will cause the model to stop generating.\n /// The Anthropic models stop on \"\\n\\nHuman:\", and may include additional built-in stop sequences in the future.\n /// By providing the stop_sequences parameter, you may include additional strings that will cause the model to stop generating.\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub stop_sequences: Option>,\n\n /// Amount of randomness injected into the response.\n ///\n /// Defaults to 1.0. Ranges from 0.0 to 1.0. Use temperature closer to 0.0 for analytical / multiple choice, and closer to 1.0 for creative and generative tasks.\n ///\n // Note that even with temperature of 0.0, the results will not be fully deterministic.\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub temperature: Option,\n\n /// Only sample from the top K options for each subsequent token.\n ///\n /// Used to remove \"long tail\" low probability responses. Learn more technical details [here](https://towardsdatascience.com/how-to-sample-from-language-models-682bceb97277).\n ///\n /// Recommended for advanced use cases only. You usually only need to use temperature.\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub top_k: Option,\n}\n\n#[derive(Debug, PartialEq, Eq)]\npub enum AnthropicMessage {\n /// A user message.\n User(String),\n /// An assistant message.\n Assistant(String),\n}\n\n#[derive(Debug)]\npub enum AnthropicMessageRole {\n User,\n Assistant,\n}\n\nimpl std::fmt::Display for AnthropicMessageRole {\n fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {\n match self {\n AnthropicMessageRole::User => write!(f, \"User\"),\n AnthropicMessageRole::Assistant => write!(f, \"Assistant\"),\n }\n }\n}\n\n#[derive(Debug, Serialize, Deserialize)]\npub struct AnthropicMessagesApiMessage {\n /// The role of the message.\n /// For a user, this will be \"user\". For an assistant, this will be \"assistant\".\n pub role: String,\n\n /// The content of the message.\n pub content: String,\n}\n\n#[derive(Debug, Serialize, Deserialize)]\n#[serde(untagged)]\npub enum AnthropicMessagesApiResponse {\n Success(AnthropicMessagesApiResponseSuccess),\n Error(AnthropicApiError),\n}\n\n/// Response from the Anthropic API for generating a response.\n/// Referenced from the Anthropic API documentation [here](https://docs.anthropic.com/en/api/complete).\n#[derive(Debug, Serialize, Deserialize)]\npub struct AnthropicMessagesApiResponseSuccess {\n /// Object type. For text completion, this will be \"completion\".\n #[serde(rename = \"type\")]\n pub typ: String,\n\n /// The role of the message. For responses, this would be \"assistant\".\n pub role: String,\n\n /// The model that generated the response.\n pub model: String,\n\n /// The resulting completion up to and excluding the stop sequences.\n pub content: Vec,\n\n /// The reason that the model stopped generating tokens.\n ///\n /// This may be one the following values:\n /// - \"stop_sequence\": Reached a stop sequence — either provided by you via the stop_sequences parameter, or a stop sequence built into the model\n /// - \"max_tokens\": Exceeded `max_tokens_to_sample` or the model's maximum\n pub stop_reason: Option,\n}\n\n/// Response from the Anthropic API for the messages API endpoint.\n#[derive(Debug, Serialize, Deserialize)]\npub struct AnthropicMessagesApiResponseSuccessContent {\n /// Type of the response, for text completion, this will be \"text\".\n #[serde(rename = \"type\")]\n pub typ: String,\n\n /// The content of the response.\n pub text: String,\n}\n\n/// Response from the Anthropic API which indicates an error.\n/// Referenced from the Anthropic API documentation [here](https://docs.anthropic.com/en/api/errors#error-shapes).\n#[derive(Debug, Serialize, Deserialize)]\npub struct AnthropicApiError {\n /// Type of the response, for error responses this will be \"error\".\n #[serde(rename = \"type\")]\n pub typ: String,\n /// Error message.\n pub error: AnthropicApiErrorBody,\n}\n\n#[derive(Debug, Serialize, Deserialize)]\npub struct AnthropicApiErrorBody {\n /// Type of the error (e.g., \"not_found_error\").\n #[serde(rename = \"type\")]\n pub typ: String,\n\n /// Error message.\n pub message: String,\n}\n" }, { "path": "core/src/lm/lm_provider/anthropic/lm.rs", "content": "use async_trait::async_trait;\nuse thiserror::Error;\n\nuse crate::lm::{\n LanguageModel, LanguageModelError, LanguageModelProvider, TextCompleteOptions,\n TextCompleteResponse, TextCompleteStreamOptions, TextCompleteStreamResponse,\n};\n\nuse super::client::{\n anthropic_client::AnthropicClientTextCompleteOptionsBuilder,\n builder::AnthropicClientBuilder,\n models::{AnthropicMessage, AnthropicMessageRole},\n};\n\n#[derive(Debug, Clone)]\npub struct Anthropic {\n pub api_endpoint: String,\n pub api_key: String,\n pub model: String,\n}\n\n#[derive(Error, Debug)]\npub enum AnthropicError {\n #[error(\"Unexpected response from API. Error: {0}\")]\n Api(String),\n\n #[error(\"Configuration error: {0}\")]\n Configuration(String),\n\n #[error(\"Serialization error: {0}\")]\n Serialization(String),\n\n #[error(\n \"OpenAi API is not available. Please check if OpenAi is running in the specified port. Error: {0}\"\n )]\n ApiUnavailable(String),\n\n #[error(\"Invalid input: {0}\")]\n InvalidInput(String),\n}\n\n#[async_trait]\nimpl LanguageModel for Anthropic {\n // TODO: Support context.\n async fn text_complete(\n &self,\n prompt: &str,\n system_prompt: &str,\n _options: TextCompleteOptions,\n ) -> Result {\n let client = AnthropicClientBuilder::new()\n .with_api_endpoint(self.api_endpoint.clone())\n .with_api_key(self.api_key.clone())\n .try_build()\n .map_err(|e| {\n LanguageModelError::Anthropic(AnthropicError::Configuration(e.to_string()))\n })?;\n\n let options = AnthropicClientTextCompleteOptionsBuilder::new()\n .with_model(self.model.clone())\n .try_build()\n .map_err(|e| {\n LanguageModelError::Anthropic(AnthropicError::Configuration(e.to_string()))\n })?;\n\n // In the case of Anthropic, we need to supply the full history of the conversation.\n // We therefore parse the prompt string and construct the messages.\n let messages = Self::messages_from_prompt(prompt)?;\n\n let response = client\n .text_complete(messages.as_slice(), system_prompt, options)\n .await\n .map_err(|e| LanguageModelError::Anthropic(AnthropicError::Api(e.to_string())))?;\n\n let response_content = response\n .content\n .first()\n .ok_or(AnthropicError::Api(\"Response content is empty\".to_string()))?;\n Ok(TextCompleteResponse {\n text: response_content.text.clone(),\n context: None,\n })\n }\n\n async fn text_complete_stream(\n &self,\n _prompt: &str,\n _system_prompt: &str,\n _options: TextCompleteStreamOptions,\n ) -> Result {\n return Err(LanguageModelError::UnsupportedFeature(\n \"Streaming is not supported for Anthropic\".to_string(),\n ));\n }\n\n async fn generate_embedding(&self, _prompt: &str) -> Result, LanguageModelError> {\n return Err(LanguageModelError::UnsupportedFeature(\n\t\t\t\t\"Embedding generation is not available on Anthropic. For more details see https://docs.anthropic.com/en/docs/build-with-claude/embeddings\".to_string(),\n\t\t\t));\n }\n\n fn provider(&self) -> LanguageModelProvider {\n LanguageModelProvider::Anthropic\n }\n\n fn text_completion_model_name(&self) -> String {\n self.model.to_string()\n }\n\n fn embedding_model_name(&self) -> String {\n \"(UNSUPPORTED)\".to_string()\n }\n}\n\nimpl Anthropic {\n fn messages_from_prompt(prompt: &str) -> Result, LanguageModelError> {\n if !prompt.starts_with(\"User:\") && !prompt.starts_with(\"Assistant:\") {\n // Assume the prompt is just the user message.\n return Ok(vec![AnthropicMessage::User(prompt.to_string())]);\n }\n\n let mut messages = Vec::new();\n\n let mut iterated_prompt = prompt.to_owned();\n\n while !iterated_prompt.trim().is_empty() {\n let current_role = if iterated_prompt.starts_with(\"User:\") {\n AnthropicMessageRole::User\n } else {\n AnthropicMessageRole::Assistant\n };\n\n let current_message = iterated_prompt\n .strip_prefix(format!(\"{}:\", current_role).as_str())\n .map(|s| s.trim())\n .ok_or(AnthropicError::InvalidInput(\n \"Prompt is not in the expected format\".to_string(),\n ))?;\n\n if !current_message.contains(\"\\n\\n\") {\n // Last message - it contains the entire content.\n match current_role {\n AnthropicMessageRole::User => {\n messages.push(AnthropicMessage::User(current_message.to_owned()));\n }\n AnthropicMessageRole::Assistant => {\n messages.push(AnthropicMessage::Assistant(current_message.to_owned()));\n }\n }\n break;\n }\n\n // Parse until the next role.\n let (current_message, next_message) =\n current_message\n .split_once(\"\\n\\n\")\n .ok_or(AnthropicError::InvalidInput(\n \"Prompt is not in the expected format\".to_string(),\n ))?;\n\n match current_role {\n AnthropicMessageRole::User => {\n messages.push(AnthropicMessage::User(current_message.to_owned()));\n }\n AnthropicMessageRole::Assistant => {\n messages.push(AnthropicMessage::Assistant(current_message.to_owned()));\n }\n }\n iterated_prompt = next_message.trim().to_string();\n }\n\n Ok(messages)\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn test_messages_from_prompt_single_message() {\n let prompt = \"Hello\";\n let messages = Anthropic::messages_from_prompt(prompt).unwrap();\n assert_eq!(messages.len(), 1);\n let Some(msg) = messages.first() else {\n panic!(\"Expected at least one message\");\n };\n let AnthropicMessage::User(content) = msg else {\n panic!(\"Expected a user message\");\n };\n assert_eq!(content, \"Hello\");\n }\n\n #[test]\n fn test_messages_from_prompt_multiple_messages() {\n let prompt = \"User: Hello\\n\\nAssistant: Hi\\n\\nUser: How are you?\";\n let messages = Anthropic::messages_from_prompt(prompt).unwrap();\n assert_eq!(messages.len(), 3);\n assert_eq!(messages[0], AnthropicMessage::User(\"Hello\".to_string()));\n assert_eq!(messages[1], AnthropicMessage::Assistant(\"Hi\".to_string()));\n assert_eq!(\n messages[2],\n AnthropicMessage::User(\"How are you?\".to_string())\n );\n }\n}\n" }, { "path": "core/src/lm/lm_provider/anthropic/mod.rs", "content": "mod builder;\nmod client;\nmod lm;\nmod models;\n\npub use builder::*;\npub use lm::*;\npub use models::*;\n" }, { "path": "core/src/lm/lm_provider/anthropic/models.rs", "content": "#![allow(dead_code)]\n\n/// Convenience constants for the Anthropic models.\npub mod anthropic_model {\n pub const CLAUDE_3_5_SONNET: &str = \"claude-3-5-sonnet-20240620\";\n pub const CLAUDE_3_OPUS: &str = \"claude-3-opus-20240229\";\n pub const CLAUDE_3_SONNET: &str = \"claude-3-sonnet-20240229\";\n pub const CLAUDE_3_HAIKU: &str = \"claude-3-haiku-20240307\";\n}\n" }, { "path": "core/src/lm/lm_provider/mod.rs", "content": "mod anthropic;\nmod models;\nmod ollama;\nmod openai;\n\npub use anthropic::*;\npub use models::*;\npub use ollama::*;\npub use openai::*;\n" }, { "path": "core/src/lm/lm_provider/models.rs", "content": "use serde::{Deserialize, Serialize};\n\nuse crate::lm::{LanguageModel, LanguageModelProviderError};\n\nuse super::{Anthropic, Ollama, OpenAi};\n\n#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]\npub enum LanguageModelProvider {\n #[serde(rename = \"ollama\")]\n Ollama,\n #[serde(rename = \"openai\")]\n OpenAi,\n #[serde(rename = \"anthropic\")]\n Anthropic,\n}\n\nimpl TryFrom<&str> for LanguageModelProvider {\n type Error = LanguageModelProviderError;\n\n fn try_from(value: &str) -> Result {\n Ok(match value {\n \"ollama\" => LanguageModelProvider::Ollama,\n \"openai\" => LanguageModelProvider::OpenAi,\n \"anthropic\" => LanguageModelProvider::Anthropic,\n _ => return Err(LanguageModelProviderError::InvalidValue(value.to_string())),\n })\n }\n}\n\nimpl LanguageModelProvider {\n /// Returns whether the provider runs local inference or not.\n pub fn is_local(&self) -> bool {\n match self {\n LanguageModelProvider::Ollama => true,\n LanguageModelProvider::OpenAi => false,\n LanguageModelProvider::Anthropic => false,\n }\n }\n}\n\npub enum OrchLanguageModel {\n Ollama(Ollama),\n OpenAi(OpenAi),\n Anthropic(Anthropic),\n}\n\nimpl OrchLanguageModel {\n pub fn provider(&self) -> LanguageModelProvider {\n match self {\n OrchLanguageModel::Ollama(_) => LanguageModelProvider::Ollama,\n OrchLanguageModel::OpenAi(_) => LanguageModelProvider::OpenAi,\n OrchLanguageModel::Anthropic(_) => LanguageModelProvider::Anthropic,\n }\n }\n\n pub fn text_completion_model_name(&self) -> String {\n match self {\n OrchLanguageModel::Ollama(lm) => lm.text_completion_model_name(),\n OrchLanguageModel::OpenAi(lm) => lm.text_completion_model_name(),\n OrchLanguageModel::Anthropic(lm) => lm.text_completion_model_name(),\n }\n }\n\n pub fn embedding_model_name(&self) -> String {\n match self {\n OrchLanguageModel::Ollama(lm) => lm.embedding_model_name(),\n OrchLanguageModel::OpenAi(lm) => lm.embedding_model_name(),\n OrchLanguageModel::Anthropic(lm) => lm.embedding_model_name(),\n }\n }\n}\n" }, { "path": "core/src/lm/lm_provider/ollama/builder.rs", "content": "use thiserror::Error;\n\nuse crate::lm::{LanguageModelBuilder, LanguageModelBuilderError};\n\nuse super::config::{DEFAULT_BASE_URL, DEFAULT_EMBEDDINGS_MODEL, DEFAULT_MODEL};\nuse super::Ollama;\n\n#[derive(Debug, Error)]\npub enum OllamaBuilderError {\n #[error(\"Configuration error: {0} is not set\")]\n ConfigurationNotSet(String),\n}\n\n/// Builds an [`Ollama`] instance.\npub struct OllamaBuilder {\n /// Base URL for the Ollama API. Defaults to [`DEFAULT_BASE_URL`].\n base_url: Option,\n /// Model to use for text completion. Defaults to [`DEFAULT_MODEL`].\n model: Option,\n /// Model to use for embedding generation. Defaults to [`DEFAULT_EMBEDDINGS_MODEL`].\n embeddings_model: Option,\n}\n\nimpl OllamaBuilder {\n /// Overrides the default base URL for the Ollama API.\n /// Defaults to [`DEFAULT_BASE_URL`].\n pub fn with_base_url(mut self, base_url: String) -> Self {\n self.base_url = Some(base_url);\n self\n }\n\n /// Overrides the default model to use for text completion.\n /// Defaults to [`DEFAULT_MODEL`].\n pub fn with_model(mut self, model: String) -> Self {\n self.model = Some(model);\n self\n }\n\n pub fn with_embeddings_model(mut self, embeddings_model: String) -> Self {\n self.embeddings_model = Some(embeddings_model);\n self\n }\n}\n\nimpl LanguageModelBuilder for OllamaBuilder {\n fn new() -> Self {\n Self {\n base_url: Some(DEFAULT_BASE_URL.to_string()),\n model: Some(DEFAULT_MODEL.to_string()),\n embeddings_model: Some(DEFAULT_EMBEDDINGS_MODEL.to_string()),\n }\n }\n\n /// Tries to build an [`Ollama`] instance. May fail if the required configurations are not set.\n fn try_build(self) -> Result {\n let Some(base_url) = self.base_url else {\n return Err(LanguageModelBuilderError::ConfigurationNotSet(\n \"Base URL\".to_string(),\n ));\n };\n let Some(model) = self.model else {\n return Err(LanguageModelBuilderError::ConfigurationNotSet(\n \"Model\".to_string(),\n ));\n };\n let Some(embeddings_model) = self.embeddings_model else {\n return Err(LanguageModelBuilderError::ConfigurationNotSet(\n \"Embeddings model\".to_string(),\n ));\n };\n Ok(Ollama {\n base_url: base_url.to_owned(),\n model: model.to_owned(),\n embeddings_model: embeddings_model.to_owned(),\n })\n }\n}\n" }, { "path": "core/src/lm/lm_provider/ollama/config.rs", "content": "use super::{ollama_embedding_model, ollama_model};\n\n/// Default base URL for the Ollama API.\npub const DEFAULT_BASE_URL: &str = \"http://localhost:11434\";\n/// Default model to use for text completion.\npub const DEFAULT_MODEL: &str = ollama_model::LLAMA3_1_8B;\n/// Default model to use for embedding generation.\npub const DEFAULT_EMBEDDINGS_MODEL: &str = ollama_embedding_model::NOMIC_EMBED_TEXT;\n" }, { "path": "core/src/lm/lm_provider/ollama/lm.rs", "content": "use async_trait::async_trait;\nuse lm::{\n error::LanguageModelError,\n models::{\n TextCompleteOptions, TextCompleteResponse, TextCompleteStreamOptions,\n TextCompleteStreamResponse,\n },\n LanguageModel, LanguageModelProvider,\n};\nuse net::SseClient;\nuse thiserror::Error;\nuse tokio_stream::StreamExt;\n\nuse crate::*;\n\nuse super::{\n OllamaApiModelsMetadata, OllamaEmbeddingsRequest, OllamaEmbeddingsResponse,\n OllamaGenerateRequest, OllamaGenerateResponse, OllamaGenerateStreamItemResponse,\n};\n\n#[derive(Debug, Clone)]\npub struct Ollama {\n pub base_url: String,\n pub model: String,\n pub embeddings_model: String,\n}\n\n#[derive(Error, Debug)]\npub enum OllamaError {\n #[error(\"Unexpected response from API. Error: {0}\")]\n Api(String),\n\n #[error(\"Unexpected error when parsing response from Ollama. Error: {0}\")]\n Parsing(String),\n\n #[error(\"{0}\")]\n Configuration(String),\n\n #[error(\"{0}\")]\n Serialization(String),\n\n #[error(\n \"Ollama API is not available. Please check if Ollama is running in the specified port. Error: {0}\"\n )]\n ApiUnavailable(String),\n}\n\nimpl Ollama {\n /// Lists the running models in the Ollama API.\n ///\n /// # Returns\n ///\n /// A [Result] containing the list of running models or an error if there was a problem.\n ///\n #[allow(dead_code)]\n pub(crate) fn list_running_models(&self) -> Result {\n let response = self.get_from_ollama_api(\"api/ps\")?;\n let parsed_response = Self::parse_models_response(&response)?;\n Ok(parsed_response)\n }\n\n // /// Lists the local models in the Ollama API.\n // ///\n // /// # Returns\n // ///\n // /// A [Result] containing the list of local models or an error if there was a problem.\n #[allow(dead_code)]\n pub fn list_local_models(&self) -> Result {\n let response = self.get_from_ollama_api(\"api/tags\")?;\n let parsed_response = Self::parse_models_response(&response)?;\n Ok(parsed_response)\n }\n\n fn parse_models_response(response: &str) -> Result {\n let models: OllamaApiModelsMetadata =\n serde_json::from_str(response).map_err(|e| OllamaError::Parsing(e.to_string()))?;\n Ok(models)\n }\n\n fn get_from_ollama_api(&self, url: &str) -> Result {\n let url = format!(\"{}/{}\", self.base_url, url);\n\n let client = reqwest::blocking::Client::new();\n let response = client\n .get(url)\n .send()\n .map_err(|e| OllamaError::ApiUnavailable(e.to_string()))?;\n let response_text = response\n .text()\n .map_err(|e| OllamaError::Api(e.to_string()))?;\n Ok(response_text)\n }\n}\n\n#[async_trait]\nimpl LanguageModel for Ollama {\n async fn text_complete(\n &self,\n prompt: &str,\n system_prompt: &str,\n _options: TextCompleteOptions,\n ) -> Result {\n let body = OllamaGenerateRequest {\n model: self.model.to_owned(),\n prompt: prompt.to_string(),\n system: Some(system_prompt.to_string()),\n ..Default::default()\n };\n\n let client = reqwest::Client::new();\n let url = format!(\"{}/api/generate\", self.base_url);\n let response = client\n .post(url)\n .body(serde_json::to_string(&body).unwrap())\n .send()\n .await\n .map_err(|e| LanguageModelError::Ollama(OllamaError::ApiUnavailable(e.to_string())))?;\n let body = response\n .text()\n .await\n .map_err(|e| LanguageModelError::Ollama(OllamaError::Api(e.to_string())))?;\n let ollama_response: OllamaGenerateResponse = serde_json::from_str(&body).map_err(|e| {\n LanguageModelError::Ollama(OllamaError::Parsing(format!(\n \"{}. Received response: {body}\",\n e\n )))\n })?;\n match ollama_response {\n OllamaGenerateResponse::Success(success_response) => Ok(TextCompleteResponse {\n text: success_response.response,\n context: success_response.context,\n }),\n OllamaGenerateResponse::Error(error_response) => Err(LanguageModelError::Ollama(\n OllamaError::Api(format!(\"{error_response:?}\")),\n )),\n }\n }\n\n async fn text_complete_stream(\n &self,\n prompt: &str,\n system_prompt: &str,\n options: TextCompleteStreamOptions,\n ) -> Result {\n let body = OllamaGenerateRequest {\n model: self.model.to_owned(),\n prompt: prompt.to_string(),\n stream: Some(true),\n format: None,\n images: None,\n system: Some(system_prompt.to_string()),\n keep_alive: Some(\"5m\".to_string()),\n context: options.context,\n };\n\n let url = format!(\"{}/api/generate\", self.base_url);\n let stream = SseClient::post(&url, Some(serde_json::to_string(&body).unwrap()));\n let stream = stream.map(|event| {\n let parsed_message = serde_json::from_str::(&event);\n match parsed_message {\n Ok(message) => match message {\n OllamaGenerateStreamItemResponse::Success(success_response) => {\n Ok(success_response.response)\n }\n OllamaGenerateStreamItemResponse::Error(error_response) => Err(\n LanguageModelError::Ollama(OllamaError::Api(format!(\"{error_response:?}\"))),\n ),\n },\n Err(e) => Err(LanguageModelError::Ollama(OllamaError::Parsing(\n e.to_string(),\n ))),\n }\n });\n let response = TextCompleteStreamResponse {\n stream: Box::pin(stream),\n };\n Ok(response)\n }\n\n async fn generate_embedding(&self, prompt: &str) -> Result, LanguageModelError> {\n let client = reqwest::Client::new();\n let url = format!(\"{}/api/embeddings\", self.base_url);\n let body = OllamaEmbeddingsRequest {\n model: self.embeddings_model.to_owned(),\n prompt: prompt.to_string(),\n };\n let response = client\n .post(url)\n .body(\n serde_json::to_string(&body)\n .map_err(|e| OllamaError::Serialization(e.to_string()))?,\n )\n .send()\n .await\n .map_err(|e| OllamaError::ApiUnavailable(e.to_string()))?;\n let body = response\n .text()\n .await\n .map_err(|e| OllamaError::Api(e.to_string()))?;\n let response: OllamaEmbeddingsResponse =\n serde_json::from_str(&body).map_err(|e| OllamaError::Parsing(e.to_string()))?;\n\n Ok(response.embedding)\n }\n\n fn provider(&self) -> LanguageModelProvider {\n LanguageModelProvider::Ollama\n }\n\n fn text_completion_model_name(&self) -> String {\n self.model.to_string()\n }\n\n fn embedding_model_name(&self) -> String {\n self.embeddings_model.to_string()\n }\n}\n" }, { "path": "core/src/lm/lm_provider/ollama/mod.rs", "content": "mod builder;\nmod config;\nmod lm;\nmod models;\n\npub use builder::*;\npub use lm::*;\npub use models::*;\n" }, { "path": "core/src/lm/lm_provider/ollama/models.rs", "content": "use serde::{Deserialize, Serialize};\n\npub mod ollama_model {\n /// https://ollama.com/library/llama3:latest\n pub const LLAMA3: &str = \"llama3:latest\";\n /// https://ollama.com/library/llama3:8b\n pub const LLAMA3_8B: &str = \"llama3:8b\";\n /// https://ollama.com/library/llama3.1:8b\n pub const LLAMA3_1_8B: &str = \"llama3.1:8b\";\n /// https://ollama.com/library/phi3:latest\n pub const PHI3_MINI: &str = \"phi3:latest\";\n /// https://ollama.com/library/codestral:latest\n pub const CODESTRAL: &str = \"codestral:latest\";\n}\n\npub mod ollama_embedding_model {\n /// https://ollama.com/library/nomic-embed-text:latest\n pub const NOMIC_EMBED_TEXT: &str = \"nomic-embed-text:latest\";\n}\n\n/// Response from the Ollama API for obtaining information about local models.\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#list-running-models).\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaApiModelsMetadata {\n pub models: Vec,\n}\n\n/// Response item from the Ollama API for obtaining information about local models.\n///\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#response-22).\n#[allow(dead_code)]\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaApiModelMetadata {\n /// The name of the model (e.g., \"mistral:latest\")\n pub name: String,\n\n /// The Ollama identifier of the model (e.g., \"mistral:latest\")\n pub model: String,\n\n /// Size of the model in bytes\n pub size: usize,\n\n /// Digest of the model using SHA256 (e.g., \"2ae6f6dd7a3dd734790bbbf58b8909a606e0e7e97e94b7604e0aa7ae4490e6d8\")\n pub digest: String,\n\n /// Model expiry time in ISO 8601 format (e.g., \"2024-06-04T14:38:31.83753-07:00\")\n pub expires_at: Option,\n\n /// More details about the model\n pub details: OllamaApiModelDetails,\n}\n\n/// Details about a running model in the API for listing running models (`GET /api/ps`).\n///\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#response-22).\n#[allow(dead_code)]\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaApiModelDetails {\n /// Model identifier that this model is based on\n pub parent_model: String,\n\n /// Format that this model is stored in (e.g., \"gguf\")\n pub format: String,\n\n /// Model family (e.g., \"ollama\")\n pub family: String,\n\n /// Parameters of the model (e.g., \"7.2B\")\n pub parameter_size: String,\n\n /// Quantization level of the model (e.g., \"Q4_0\" for 4-bit quantization)\n pub quantization_level: String,\n}\n\n/// Request for generating a response from the Ollama API.\n///\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#generate-a-completion).\n#[allow(dead_code)]\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaGenerateRequest {\n /// Model identifier (e.g., \"mistral:latest\")\n pub model: String,\n\n /// The prompt to generate a response for (e.g., \"List all Kubernetes pods\")\n pub prompt: String,\n\n /// The context parameter returned from a previous request to /generate, this can be used to keep a short conversational memory\n pub context: Option>,\n\n /// Optional list of base64-encoded images (for multimodal models such as `llava`)\n pub images: Option>,\n\n /// Optional format to use for the response (currently only \"json\" is supported)\n pub format: Option,\n\n /// Optional flag that controls whether the response is streamed or not (defaults to true).\n /// If `false`` the response will be returned as a single response object, rather than a stream of objects\n pub stream: Option,\n\n // System message (overrides what is defined in the Modelfile)\n pub system: Option,\n\n /// Controls how long the model will stay loaded into memory following the request (default: 5m)\n pub keep_alive: Option,\n}\n\nimpl Default for OllamaGenerateRequest {\n fn default() -> Self {\n Self {\n model: ollama_model::CODESTRAL.to_string(),\n prompt: \"\".to_string(),\n stream: Some(false),\n format: None,\n images: None,\n system: Some(\"You are a helpful assistant\".to_string()),\n keep_alive: Some(\"5m\".to_string()),\n context: None,\n }\n }\n}\n\n/// Response from the Ollama API for generating a response.\n#[derive(Debug, Serialize, Deserialize)]\n#[serde(untagged)]\npub enum OllamaGenerateResponse {\n Success(OllamaGenerateResponseSuccess),\n Error(OllamaGenerateResponseError),\n}\n\n#[derive(Debug, Serialize, Deserialize)]\n#[allow(dead_code)]\npub struct OllamaGenerateResponseSuccess {\n /// Model identifier (e.g., \"mistral:latest\")\n pub model: String,\n\n /// Time at which the response was generated (ISO 8601 format)\n pub created_at: String,\n\n /// The response to the prompt\n pub response: String,\n\n /// The encoding of the conversation used in this response, this can be sent in the next request to keep a conversational memory\n pub context: Option>,\n\n /// The duration of the response in nanoseconds\n pub total_duration: usize,\n}\n\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaGenerateResponseError {\n /// Error message.\n pub error: String,\n}\n\n#[derive(Debug, Serialize, Deserialize)]\n#[serde(untagged)]\npub enum OllamaGenerateStreamItemResponse {\n Success(OllamaGenerateStreamItemResponseSuccess),\n Error(OllamaGenerateStreamItemResponseError),\n}\n\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaGenerateStreamItemResponseSuccess {\n /// Model identifier (e.g., \"mistral:latest\")\n pub model: String,\n\n /// Time at which the response was generated (ISO 8601 format)\n pub created_at: String,\n\n /// The response to the prompt\n pub response: String,\n}\n\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaGenerateStreamItemResponseError {\n /// Error message.\n pub error: String,\n}\n\n/// Request for generating an embedding from the Ollama API.\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#generate-embeddings).\n///\n#[allow(dead_code)]\n#[derive(Debug, Serialize, Deserialize, Default)]\npub struct OllamaEmbeddingsRequest {\n /// The string to generate an embedding for.\n pub prompt: String,\n\n /// The model to use for the embedding generation.\n pub model: String,\n}\n\n/// Response from the Ollama API for generating an embedding.\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#generate-embeddings).\n///\n#[allow(dead_code)]\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaEmbeddingsResponse {\n /// The embedding for the prompt.\n pub embedding: Vec,\n}\n" }, { "path": "core/src/lm/lm_provider/openai/builder.rs", "content": "use thiserror::Error;\n\nuse crate::lm::{lm_provider::openai::config, LanguageModelBuilder, LanguageModelBuilderError};\n\nuse super::OpenAi;\n\n#[derive(Debug, Error)]\npub enum OpenAiBuilderError {\n #[error(\"Configuration error: {0} is not set\")]\n ConfigurationNotSet(String),\n}\n\n/// Builds an [`OpenAi`] instance.\npub struct OpenAiBuilder {\n api_endpoint: Option,\n api_key: Option,\n model: Option,\n embeddings_model: Option,\n}\n\nimpl OpenAiBuilder {\n /// Sets the required API key for the OpenAI API.\n pub fn with_api_key(mut self, api_key: String) -> Self {\n self.api_key = Some(api_key);\n self\n }\n\n /// Overrides the default API endpoint for the OpenAI API.\n pub fn with_api_endpoint(mut self, api_endpoint: String) -> Self {\n self.api_endpoint = Some(api_endpoint);\n self\n }\n\n /// Sets the model to use for text completion. Defaults to [`config::DEFAULT_MODEL`].\n pub fn with_model(mut self, model: String) -> Self {\n self.model = Some(model);\n self\n }\n\n /// Sets the model to use for embedding generation. Defaults to [`config::DEFAULT_EMBEDDINGS_MODEL`].\n pub fn with_embeddings_model(mut self, embeddings_model: String) -> Self {\n self.embeddings_model = Some(embeddings_model.clone());\n self\n }\n}\n\nimpl LanguageModelBuilder for OpenAiBuilder {\n fn new() -> Self {\n Self {\n api_key: None,\n api_endpoint: None,\n model: Some(config::DEFAULT_MODEL.to_string()),\n embeddings_model: Some(config::DEFAULT_EMBEDDINGS_MODEL.to_string()),\n }\n }\n\n /// Tries to build an [`OpenAi`] instance. May fail if the required configurations are not set.\n fn try_build(self) -> Result {\n let Some(api_key) = self.api_key else {\n return Err(LanguageModelBuilderError::ConfigurationNotSet(\n \"API key\".to_string(),\n ));\n };\n let Some(model) = self.model else {\n return Err(LanguageModelBuilderError::ConfigurationNotSet(\n \"Model\".to_string(),\n ));\n };\n let Some(embeddings_model) = self.embeddings_model else {\n return Err(LanguageModelBuilderError::ConfigurationNotSet(\n \"Embeddings model\".to_string(),\n ));\n };\n Ok(OpenAi {\n api_endpoint: self.api_endpoint,\n api_key: api_key.to_owned(),\n model: model.to_owned(),\n embeddings_model: embeddings_model.to_owned(),\n })\n }\n}\n" }, { "path": "core/src/lm/lm_provider/openai/config.rs", "content": "use super::{openai_embedding_model, openai_model};\n\n/// Default model to use for text completion.\npub const DEFAULT_MODEL: &str = openai_model::GPT_4O_MINI;\n/// Default model to use for embedding generation.\npub const DEFAULT_EMBEDDINGS_MODEL: &str = openai_embedding_model::TEXT_EMBEDDING_ADA_002;\n" }, { "path": "core/src/lm/lm_provider/openai/lm.rs", "content": "use async_trait::async_trait;\nuse lm::{\n error::LanguageModelError,\n models::{\n TextCompleteOptions, TextCompleteResponse, TextCompleteStreamOptions,\n TextCompleteStreamResponse,\n },\n LanguageModel, LanguageModelProvider,\n};\nuse openai_api_rs::v1::{\n api::OpenAIClient,\n chat_completion::{self, ChatCompletionRequest},\n embedding::EmbeddingRequest,\n};\nuse thiserror::Error;\nuse tokio_stream::{self as stream};\n\nuse crate::*;\n\n#[derive(Debug, Clone)]\npub struct OpenAi {\n pub api_endpoint: Option,\n pub api_key: String,\n pub model: String,\n pub embeddings_model: String,\n}\n\n#[derive(Error, Debug)]\npub enum OpenAiError {\n #[error(\"Unexpected response from API. Error: {0}\")]\n Api(String),\n\n #[error(\"Configuration error: {0}\")]\n Configuration(String),\n\n #[error(\"Serialization error: {0}\")]\n Serialization(String),\n\n #[error(\"OpenAI API is not available. Error: {0}\")]\n ApiUnavailable(String),\n}\n\n#[async_trait]\nimpl LanguageModel for OpenAi {\n async fn text_complete(\n &self,\n prompt: &str,\n system_prompt: &str,\n _options: TextCompleteOptions,\n ) -> Result {\n let mut client = OpenAIClient::new(self.api_key.to_owned());\n\n if let Some(api_endpoint) = self.api_endpoint.clone() {\n client.api_endpoint = api_endpoint;\n }\n\n let messages = vec![\n chat_completion::ChatCompletionMessage {\n role: chat_completion::MessageRole::system,\n content: chat_completion::Content::Text(system_prompt.to_owned()),\n name: None,\n tool_calls: None,\n tool_call_id: None,\n },\n chat_completion::ChatCompletionMessage {\n role: chat_completion::MessageRole::user,\n content: chat_completion::Content::Text(prompt.to_owned()),\n name: None,\n tool_calls: None,\n tool_call_id: None,\n },\n ];\n // TODO: Support customization of max tokens and temperature.\n let req = ChatCompletionRequest::new(self.model.to_owned(), messages);\n\n let result = client\n .chat_completion(req)\n .await\n .map_err(|e| LanguageModelError::OpenAi(OpenAiError::Api(e.to_string())))?;\n let completion = result\n .choices\n .first()\n .unwrap()\n .message\n .content\n .clone()\n .unwrap();\n Ok(TextCompleteResponse {\n text: completion,\n // TODO: Support context.\n context: None,\n })\n }\n\n async fn text_complete_stream(\n &self,\n prompt: &str,\n system_prompt: &str,\n _options: TextCompleteStreamOptions,\n ) -> Result {\n // TODO: Support streaming - currently it just sends a single message.\n let text_completion_response = self\n .text_complete(prompt, system_prompt, TextCompleteOptions { context: None })\n .await?;\n Ok(TextCompleteStreamResponse {\n stream: Box::pin(stream::once(Ok(text_completion_response.text))),\n })\n }\n\n async fn generate_embedding(&self, prompt: &str) -> Result, LanguageModelError> {\n let client = OpenAIClient::new(self.api_key.to_owned());\n\n let resp = client\n .embedding(EmbeddingRequest {\n model: self.embeddings_model.to_owned(),\n input: prompt.to_owned(),\n dimensions: None,\n user: None,\n })\n .await\n .map_err(|e| LanguageModelError::OpenAi(OpenAiError::Api(e.to_string())))?;\n\n let data = resp.data.first().expect(\"Embedding data not found\");\n Ok(data.embedding.clone())\n }\n\n fn provider(&self) -> LanguageModelProvider {\n LanguageModelProvider::OpenAi\n }\n\n fn text_completion_model_name(&self) -> String {\n self.model.to_string()\n }\n\n fn embedding_model_name(&self) -> String {\n self.embeddings_model.to_string()\n }\n}\n" }, { "path": "core/src/lm/lm_provider/openai/mod.rs", "content": "mod builder;\nmod config;\nmod lm;\nmod models;\n\npub use builder::*;\npub use lm::*;\npub use models::*;\n" }, { "path": "core/src/lm/lm_provider/openai/models.rs", "content": "pub mod openai_model {\n pub const GPT_3_5_TURBO: &str = \"gpt-3.5-turbo\";\n pub const GPT_4: &str = \"gpt-4\";\n pub const GPT_4O_TURBO: &str = \"gpt-4o-turbo\";\n pub const GPT_4O_MINI: &str = \"gpt-4o-mini\";\n}\n\npub mod openai_embedding_model {\n pub const TEXT_EMBEDDING_ADA_002: &str = \"text-embedding-ada-002\";\n pub const TEXT_EMBEDDING_ADA_002_DIMENSIONS: usize = 1536;\n pub const TEXT_EMBEDDING_3_SMALL: &str = \"text-embedding-3-small\";\n pub const TEXT_EMBEDDING_3_SMALL_DIMENSIONS: usize = 1536;\n pub const TEXT_EMBEDDING_3_LARGE: &str = \"text-embedding-3-large\";\n pub const TEXT_EMBEDDING_3_LARGE_DIMENSIONS: usize = 3072;\n}\n" }, { "path": "core/src/lm/mod.rs", "content": "//! A module containing all logic related to LMs (Language Models).\n//! This don't strictly have to be *large* language models (i.e., SLMs such as Phi-3 or Mistral NeMo are included).\n\nmod builder;\nmod error;\nmod lm_provider;\nmod models;\n\npub use builder::*;\npub use error::*;\npub use lm_provider::*;\npub use models::*;\n" }, { "path": "core/src/lm/models.rs", "content": "#![allow(dead_code)]\n\nuse std::pin::Pin;\n\nuse async_trait::async_trait;\nuse dyn_clone::DynClone;\nuse tokio_stream::Stream;\n\nuse super::{error::LanguageModelError, LanguageModelProvider};\n\n/// A trait for language model providers which implements text completion, embeddings, etc.\n///\n/// > `DynClone` is used so that there can be dynamic dispatch of the `Llm` trait,\n/// > especially needed for [magic-cli](https://github.com/guywaldman/magic-cli).\n#[async_trait]\npub trait LanguageModel: DynClone + Send + Sync {\n /// Generates a response from the LLM.\n ///\n /// # Arguments\n /// * `prompt` - The prompt to generate a response for.\n /// * `system_prompt` - The system prompt to use for the generation.\n /// * `options` - The options for the generation.\n ///\n /// # Returns\n /// A [Result] containing the response from the LLM or an error if there was a problem.\n ///\n async fn text_complete(\n &self,\n prompt: &str,\n system_prompt: &str,\n options: TextCompleteOptions,\n ) -> Result;\n\n /// Generates a streaming response from the LLM.\n ///\n /// # Arguments\n /// * `prompt` - The prompt to generate a response for.\n /// * `system_prompt` - The system prompt to use for the generation.\n /// * `options` - The options for the generation.\n ///\n /// # Returns\n /// A [Result] containing the response from the LLM or an error if there was a problem.\n ///\n async fn text_complete_stream(\n &self,\n prompt: &str,\n system_prompt: &str,\n options: TextCompleteStreamOptions,\n ) -> Result;\n\n /// Generates an embedding from the LLM.\n ///\n /// # Arguments\n /// * `prompt` - The item to generate an embedding for.\n ///\n /// # Returns\n ///\n /// A [Result] containing the embedding or an error if there was a problem.\n async fn generate_embedding(&self, prompt: &str) -> Result, LanguageModelError>;\n\n /// Returns the provider of the LLM.\n fn provider(&self) -> LanguageModelProvider;\n\n /// Returns the name of the model used for text completions.\n fn text_completion_model_name(&self) -> String;\n\n /// Returns the name of the model used for embeddings.\n fn embedding_model_name(&self) -> String;\n}\n\n#[derive(Debug, Clone, Default)]\npub struct TextCompleteOptions {\n /// An encoding of the conversation used in this response, this can be sent in the next request to keep a conversational memory.\n /// This should be as returned from the previous response.\n pub context: Option>,\n}\n\n#[derive(Debug, Clone, Default)]\npub struct TextCompleteStreamOptions {\n pub context: Option>,\n}\n\n#[derive(Debug, Clone)]\npub struct TextCompleteResponse {\n pub text: String,\n // TODO: This is specific to Ollama, context looks differently for other LLM providers.\n pub context: Option>,\n}\n\npub struct TextCompleteStreamResponse {\n pub stream: Pin> + Send>>,\n // TODO: Handle context with streaming response.\n // pub context: Vec,\n}\n" }, { "path": "core/src/net/mod.rs", "content": "/// Module for working with Server-Sent Events.\nmod sse;\n\npub use sse::*;\n" }, { "path": "core/src/net/sse.rs", "content": "use async_gen::AsyncIter;\nuse reqwest::{header, Client};\nuse tokio_stream::Stream;\n\n/// A client for working with Server-Sent Events.\npub struct SseClient;\n\nimpl SseClient {\n pub fn post(url: &str, body: Option) -> impl Stream {\n let client = Client::new();\n let mut req = Client::post(&client, url)\n .header(header::ACCEPT, \"text/event-stream\")\n .header(header::CACHE_CONTROL, \"no-cache\")\n .header(header::CONNECTION, \"keep-alive\")\n .header(header::CONTENT_TYPE, \"application/json\");\n if let Some(body) = body {\n req = req.body(body);\n }\n let req = req.build().unwrap();\n\n AsyncIter::from(async_gen::gen! {\n let mut conn = client.execute(req).await.unwrap();\n while let Some(event) = conn.chunk().await.unwrap() {\n yield std::str::from_utf8(&event).unwrap().to_owned();\n }\n })\n }\n}\n" }, { "path": "core/src/response.rs", "content": "pub use orch_response::*;\npub use orch_response_derive::*;\n" }, { "path": "orch.code-workspace", "content": "{\n\t\"folders\": [\n\t\t{\n\t\t\t\"path\": \"core\"\n\t\t},\n\t\t{\n\t\t\t\"path\": \"response\"\n\t\t},\n\t\t{\n\t\t\t\"path\": \"response_derive\"\n\t\t},\n\t\t{\n\t\t\t\"path\": \".\",\n\t\t\t\"name\": \"(root)\"\n\t\t}\n\t],\n\t\"settings\": {}\n}" }, { "path": "response/.gitignore", "content": "" }, { "path": "response/Cargo.toml", "content": "[package]\nname = \"orch_response\"\nversion = \"0.0.16\"\nedition = \"2021\"\nlicense = \"MIT\"\ndescription = \"Models for orch Executor responses\"\nhomepage = \"https://github.com/guywaldman/orch\"\nrepository = \"https://github.com/guywaldman/orch\"\nkeywords = [\"llm\", \"openai\", \"ollama\", \"rust\"]\n\n[dependencies]\ndyn-clone = \"1.0.17\"\nserde = \"1.0.204\"\nserde_json = \"1.0.120\"\n" }, { "path": "response/src/lib.rs", "content": "use dyn_clone::DynClone;\n\n/// Represents an option for the response of a language model.\n#[derive(Debug, Clone)]\npub struct ResponseOption {\n /// The discriminator for the response type (e.g., `Answer` or `Fail`).\n pub type_name: String,\n /// The scenario for the response (e.g., \"You know the answer\" or \"You don't know the answer\").\n pub scenario: String,\n /// The description of what the response represents (e.g., \"The capital city of the received country\" or \"Explanation on why the capital city is not known\").\n pub description: String,\n /// The schema for the response.\n pub schema: Vec,\n}\n\n/// Represents a field in the schema of a response.\n#[derive(Debug, Clone)]\npub struct ResponseSchemaField {\n /// Name of the field (e.g., \"capital\" for the capital city).\n pub name: String,\n /// Description of the field (e.g., \"Capital city of the received country\").\n pub description: String,\n /// Type of the field (e.g., \"string\" for a string).\n pub typ: String,\n /// Example of the field (e.g., \"London\" for the capital city).\n pub example: String,\n}\n\npub trait OrchResponseVariant: Send + Sync {\n fn variant() -> ResponseOption;\n}\n\npub trait OrchResponseVariants: DynClone + Send + Sync {\n fn variants(&self) -> Vec;\n fn parse(&self, response: &str) -> Result;\n}\n" }, { "path": "response_derive/.rustfmt.toml", "content": "max_width = 140\n" }, { "path": "response_derive/Cargo.toml", "content": "[package]\nname = \"orch_response_derive\"\nversion = \"0.0.16\"\nedition = \"2021\"\nlicense = \"MIT\"\ndescription = \"Derive macros for orch Executor responses\"\nhomepage = \"https://github.com/guywaldman/orch\"\nrepository = \"https://github.com/guywaldman/orch\"\nkeywords = [\"llm\", \"openai\", \"ollama\", \"rust\"]\n\n[lib]\nproc-macro = true\n\n[dependencies]\norch_response = { path = \"../response\", version = \"0.0.16\" }\ndarling = \"0.20.10\"\nproc-macro2 = \"1.0.86\"\nquote = \"1.0.36\"\nsyn = \"2.0.71\"\n\n[dev-dependencies]\nserde = { version = \"1.0.204\", features = [\"derive\"] }\nserde_json = \"1.0.120\"\n" }, { "path": "response_derive/README.md", "content": "# Procedural Macros for Orch\n\n## Resources\n\n- https://www.freecodecamp.org/news/procedural-macros-in-rust/" }, { "path": "response_derive/src/attribute_impl.rs", "content": "use darling::FromMeta;\n\n/// #[variant(...)]\n#[derive(Debug, FromMeta)]\npub(crate) struct VariantAttribute {\n pub(crate) variant: String,\n pub(crate) scenario: String,\n pub(crate) description: String,\n}\n\n/// #[schema(...)]\n#[derive(Debug, FromMeta)]\npub(crate) struct SchemaAttribute {\n pub(crate) description: String,\n pub(crate) example: String,\n}\n" }, { "path": "response_derive/src/derive_impl.rs", "content": "use darling::FromMeta;\nuse quote::quote;\nuse syn::{parse_macro_input, spanned::Spanned, DeriveInput, PathArguments};\n\nuse crate::attribute_impl::{SchemaAttribute, VariantAttribute};\n\npub(crate) fn response_variants_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {\n let mut output = quote!();\n\n // Bring traits into scope.\n output.extend(quote! {\n use ::orch_response::OrchResponseVariant;\n use ::serde::de::Error;\n });\n\n let original_enum = parse_macro_input!(input as DeriveInput);\n let DeriveInput { data, ident, .. } = original_enum.clone();\n let syn::Data::Enum(data) = data else {\n panic!(\"#[derive(OrchResponseVariants)] can only be used with enums\");\n };\n let original_enum_ident = ident;\n\n let vec_capacity = data.variants.len();\n\n let mut options_vec_pushes = quote!();\n for variant in data.variants.iter() {\n let ident = syn::Ident::new(\n &get_enum_variant_struct_ident(variant).expect(\"Failed to parse enum variant\"),\n variant.ident.span(),\n );\n\n options_vec_pushes.extend(quote! {\n options.push(#ident::variant());\n });\n }\n\n // We construct a new struct that will be used to parse the response.\n // NOTE: This is hacky, but a workaround for the fact that the enum cannot be constructed.\n let derived_enum_struct_ident = syn::Ident::new(&format!(\"{}Derived\", original_enum_ident), original_enum_ident.span());\n\n output.extend(quote! {\n #[derive(::std::fmt::Debug, ::std::clone::Clone)]\n pub struct #derived_enum_struct_ident;\n });\n\n // Note: We parse with a dynamic evaluation and looking at the `response_type` field, but this could be done\n // by deriving #[serde(tag = \"response_type\")] on the enum.\n let mut response_type_arms = quote!();\n for variant in data.variants.iter() {\n let variant_ident = variant.ident.clone();\n let variant_ident_str = syn::LitStr::new(&variant.ident.to_string(), variant.ident.span());\n let struct_ident = syn::Ident::new(\n &get_enum_variant_struct_ident(variant).expect(\"Failed to parse enum variant\"),\n variant.ident.span(),\n );\n response_type_arms.extend(quote! {\n #variant_ident_str => Ok(#original_enum_ident::#variant_ident(serde_json::from_str::<#struct_ident>(response)?)),\n });\n }\n\n output.extend(quote! {\n impl ::orch_response::OrchResponseVariants<#original_enum_ident> for #derived_enum_struct_ident {\n fn variants(&self) -> Vec<::orch_response::ResponseOption> {\n let mut options = Vec::with_capacity(#vec_capacity);\n #options_vec_pushes\n options\n }\n\n fn parse(&self, response: &str) -> Result<#original_enum_ident, ::serde_json::Error> {\n let dynamic_parsed = serde_json::from_str::(response)?;\n let response_type = dynamic_parsed.get(\"response_type\");\n let response_type = match response_type {\n Some(response_type) => match response_type.as_str() {\n Some(response_type) => response_type,\n None => {\n return Err(::serde_json::Error::custom(format!(\n \"Invalid response type: {}\",\n response\n )));\n }\n }\n None => {\n return Err(::serde_json::Error::custom(format!(\n \"Invalid response type: {}\",\n response\n )));\n }\n };\n match response_type {\n #response_type_arms\n _ => Err(::serde_json::Error::custom(\"Invalid response type\")),\n }\n }\n }\n });\n\n output.into()\n}\n\npub fn response_variant_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {\n let original_struct = parse_macro_input!(input as DeriveInput);\n let DeriveInput { data, ident, attrs, .. } = original_struct.clone();\n let syn::Data::Struct(data) = data else {\n panic!(\"#[derive(OrchResponseOption)] can only be used with structs\");\n };\n let original_struct_ident = ident.clone();\n\n let fields = data.fields;\n\n // Parse the #[variant(...)] attribute.\n let variant_attr = attrs\n .iter()\n .filter_map(|attr| VariantAttribute::from_meta(&attr.meta).ok())\n .next()\n .expect(\"#[variant(...)] attribute not found on variant field\");\n let VariantAttribute {\n variant,\n scenario,\n description,\n } = variant_attr;\n\n // Parse the fields used in [`orch::response::OrchResponseVariant`].\n let mut schema_fields = Vec::new();\n for variant_field in fields.iter() {\n // Parse the #[schema(...)] attribute.\n let schema_attr = variant_field\n .attrs\n .iter()\n .filter_map(|attr| SchemaAttribute::from_meta(&attr.meta).ok())\n .collect::>();\n if schema_attr.len() != 1 {\n panic!(\"Expected a single #[schema(...)] attribute for each field of the enum variant with the correct format and parameters\");\n }\n let SchemaAttribute { description, example } = schema_attr.first().expect(\"Failed to parse schema attribute\");\n\n let typ = ast_type_to_str(&variant_field.ty).unwrap_or_else(|_| {\n panic!(\n \"Failed to convert type to string for field `{}` of variant `{}`\",\n variant_field.ident.as_ref().unwrap(),\n ident\n )\n });\n let typ = syn::LitStr::new(&typ, variant_field.span());\n let field_ident = syn::LitStr::new(&variant_field.ident.as_ref().unwrap().to_string(), variant_field.span());\n schema_fields.push(quote! {\n ::orch_response::ResponseSchemaField {\n name: #field_ident.to_string(),\n description: #description.to_string(),\n typ: #typ.to_string(),\n example: #example.to_string(),\n }\n })\n }\n\n quote! {\n impl ::orch_response::OrchResponseVariant for #original_struct_ident {\n fn variant() -> ::orch_response::ResponseOption {\n ::orch_response::ResponseOption {\n type_name: #variant.to_string(),\n scenario: #scenario.to_string(),\n description: #description.to_string(),\n schema: vec![\n #(#schema_fields),*\n ]\n }\n }\n }\n }\n .into()\n}\n\n// Parse `Answer(AnswerResponseOption)` into `AnswerResponseOption`.\nfn get_enum_variant_struct_ident(variant: &syn::Variant) -> Result {\n // We expect the enum variant to look like this: `Answer(AnswerResponseOption)`,\n // so we parse the `AnswerResponseOption` struct.\n let syn::Fields::Unnamed(fields) = &variant.fields else {\n panic!(\"Expected an unnamed struct for each enum variant\");\n };\n let Some(syn::Field { ty, .. }) = fields.unnamed.first() else {\n panic!(\"Expected an unnamed struct for each enum variant\");\n };\n let syn::Type::Path(p) = ty else {\n panic!(\"Expected an unnamed struct for each enum variant\");\n };\n let ident = &p.path.segments.first().unwrap().ident;\n Ok(ident.to_string())\n}\n\nfn ast_type_to_str(ty: &syn::Type) -> Result {\n match ty {\n syn::Type::Path(tp) => {\n let ps = tp.path.segments.first();\n let Some(first_path_segment) = ps else {\n return Err(format!(\"Unsupported/unexpected type: {:?}\", ty).to_owned());\n };\n let t = first_path_segment.ident.to_string();\n match t.as_ref() {\n \"String\" => {\n // SUPPORTED: String\n Ok(\"string\".to_owned())\n }\n \"bool\" => {\n // SUPPORTED: bool\n Ok(\"boolean\".to_owned())\n }\n \"Option\" => {\n let PathArguments::AngleBracketed(ab) = &tp.path.segments.first().unwrap().arguments else {\n return Err(format!(\"Unsupported/unexpected type: {:?}\", ty).to_owned());\n };\n let syn::GenericArgument::Type(t) = ab.args.first().unwrap() else {\n return Err(format!(\"Unsupported/unexpected type: {:?}\", ty).to_owned());\n };\n let syn::Type::Path(p) = t else {\n return Err(format!(\"Unsupported/unexpected type: {:?}\", ty).to_owned());\n };\n let t = p.path.segments.first().unwrap().ident.to_string();\n\n match t.as_ref() {\n \"String\" => {\n // SUPPORTED: Option\n Ok(\"string?\".to_owned())\n }\n \"bool\" => {\n // SUPPORTED: Option\n Ok(\"boolean?\".to_owned())\n }\n _ => Err(format!(\"Unsupported/unexpected type: {}\", t).to_owned()),\n }\n }\n \"Vec\" => {\n let PathArguments::AngleBracketed(ab) = &tp.path.segments.first().unwrap().arguments else {\n return Err(format!(\"Unsupported/unexpected type: {:?}\", ty).to_owned());\n };\n let syn::GenericArgument::Type(t) = ab.args.first().unwrap() else {\n return Err(format!(\"Unsupported/unexpected type: {:?}\", ty).to_owned());\n };\n let syn::Type::Path(p) = t else {\n return Err(format!(\"Unsupported/unexpected type: {:?}\", ty).to_owned());\n };\n let t = p.path.segments.first().unwrap().ident.to_string();\n match t.as_ref() {\n // SUPPORTED: Vec\n \"String\" => Ok(\"string[]\".to_owned()),\n _ => Err(format!(\"Unsupported/unexpected type: {}\", t).to_owned()),\n }\n }\n _ => Err(format!(\"Unsupported/unexpected type: {}\", t).to_owned()),\n }\n }\n _ => Err(format!(\"Unsupported/unexpected type: {:?}\", ty).to_owned()),\n }\n}\n" }, { "path": "response_derive/src/lib.rs", "content": "mod attribute_impl;\nmod derive_impl;\n\nuse quote::quote;\n\n/// Used to derive the `OrchResponseVariants` trait for a given enum\n#[proc_macro_derive(Variants)]\npub fn derive_orch_response_variants(input: proc_macro::TokenStream) -> proc_macro::TokenStream {\n derive_impl::response_variants_derive(input)\n}\n\n/// Used to derive the `OrchResponseVariant` trait for a given enum.\n#[proc_macro_derive(Variant, attributes(variant, schema))]\npub fn derive_orch_response_variant_variant(input: proc_macro::TokenStream) -> proc_macro::TokenStream {\n derive_impl::response_variant_derive(input)\n}\n\n/// Used to construct the identifier of the derived enum.\n#[proc_macro]\npub fn variants(input: proc_macro::TokenStream) -> proc_macro::TokenStream {\n // Expects the identifier of the derived enum.\n let enum_ident = syn::parse_macro_input!(input as syn::Ident);\n let derived_enum_ident = syn::Ident::new(&format!(\"{}Derived\", enum_ident), enum_ident.span());\n quote! {\n #derived_enum_ident {}\n }\n .into()\n}\n" }, { "path": "scripts/ci.sh", "content": "#!/bin/bash\n\nset -e\n\nfor i in 1 2 3; do\n\tsystemctl is-active ollama.service && sudo systemctl stop ollama.service\n\tcurl -fsSL https://ollama.com/install.sh | sh\n\tsleep 5\n\tif systemctl is-active ollama.service; then\n\t\tbreak\n\tfi\ndone\n\nollama serve &\nollama pull phi3:mini\nollama pull nomic-embed-text:latest\n" }, { "path": "scripts/examples.sh", "content": "#!/usr/bin/env bash\n\nSCRIPT_DIR=\"$(cd -- \"$(dirname -- \"${BASH_SOURCE[0]}\")\" &>/dev/null && pwd)\"\nsource \"$SCRIPT_DIR/utils.sh\"\n\nprovider=\"$1\"\n# If provider is not provided, or is not one of the supported providers, exit with an error\nif [[ \"$provider\" != \"ollama\" && \"$provider\" != \"openai\" && \"$provider\" != \"anthropic\" ]]; then\n\techo \"ERROR: Invalid provider. Supported providers: ollama, openai, anthropic\"\n\texit 1\nfi\n\nFAILURE=0\n\ninfo \"Running all examples in the 'orch' crate for provider $provider...\"\npushd core 2>&1 >/dev/null\nfor example in $(find examples -name '*.rs'); do\n\texample=${example%.rs}\n\tinfo \"Running example: $(basename $example)\"\n\tcargo run --quiet --example $(basename $example) -- $provider 1>/dev/null\n\tif [ $? -ne 0 ]; then\n\t\tFAILURE=1\n\t\terror \"Example $(basename $example) failed\"\n\tfi\n\tsuccess \"Example $(basename $example) succeeded\"\ndone\npopd 2>&1 >/dev/null\ninfo \"Ran all examples in the 'orch' crate for provider $provider\"\n\nexit $FAILURE\n" }, { "path": "scripts/utils.sh", "content": "#!/usr/bin/env bash\n\nif [ \"${CI:-false}\" != \"true\" ]; then\n\tbold=$(tput bold)\n\tnormal=$(tput sgr0)\n\tlight_grey=$(tput setaf 250)\n\tblue=$(tput setaf 4)\n\tgreen=$(tput setaf 2)\n\tyellow=$(tput setaf 3)\n\tred=$(tput setaf 1)\nelse\n\tbold=\"\"\n\tnormal=\"\"\n\tblue=\"\"\n\tlight_grey=\"\"\n\tgreen=\"\"\n\tyellow=\"\"\n\tred=\"\"\nfi\n\nfunction formatted_time {\n\tdate +%FT%T.%3N\n}\n\nfunction formatted_severity {\n\tprintf \"%+6s:\" $1\n}\n\nfunction formatted_log {\n\tlog_severity=$1\n\tlog_message=\"$2\"\n\techo \"${bold}$(formatted_time) $(formatted_severity $log_severity) $log_message ${normal}\"\n}\n\nfunction info {\n\tformatted_log INFO \"$1\"\n}\n\nfunction warn {\n\tformatted_log WARN \"$1\"\n}\n\nfunction success {\n\techo \"${green}$(formatted_log INFO \"$1\")\"\n}\n\nfunction error {\n\techo >&2 \"${red}$(formatted_log ERROR \"$1\")\"\n}\n\nfunction error-and-exit {\n\terror \"$1\"\n\tif [ \"${2:''}\" != \"\" ]; then\n\t\texit $2\n\telse\n\t\texit 1\n\tfi\n}\n" }, { "path": "src/core/mod.rs", "content": "mod net;\n\npub use net::*;\n" }, { "path": "src/core/net/mod.rs", "content": "/// Module for working with Server-Sent Events.\nmod sse;\n\npub use sse::*;\n" }, { "path": "src/core/net/sse.rs", "content": "use async_gen::AsyncIter;\nuse reqwest::{header, Client};\nuse tokio_stream::Stream;\n\n/// A client for working with Server-Sent Events.\npub struct SseClient;\n\nimpl SseClient {\n pub fn post(url: &str, body: Option) -> impl Stream {\n let client = Client::new();\n let mut req = Client::post(&client, url)\n .header(header::ACCEPT, \"text/event-stream\")\n .header(header::CACHE_CONTROL, \"no-cache\")\n .header(header::CONNECTION, \"keep-alive\")\n .header(header::CONTENT_TYPE, \"application/json\");\n if let Some(body) = body {\n req = req.body(body);\n }\n let req = req.build().unwrap();\n\n AsyncIter::from(async_gen::gen! {\n let mut conn = client.execute(req).await.unwrap();\n while let Some(event) = conn.chunk().await.unwrap() {\n yield std::str::from_utf8(&event).unwrap().to_owned();\n }\n })\n }\n}\n" }, { "path": "src/executor.rs", "content": "use std::pin::Pin;\n\nuse thiserror::Error;\nuse tokio_stream::Stream;\n\nuse crate::{Llm, LlmError, TextCompleteOptions, TextCompleteStreamOptions};\n\npub struct Executor<'a, L: Llm> {\n llm: &'a L,\n}\n\n#[derive(Debug, Error)]\npub enum ExecutorError {\n #[error(\"LLM error: {0}\")]\n Llm(LlmError),\n}\n\nimpl<'a, L: Llm> Executor<'a, L> {\n /// Creates a new `Executor` instance.\n ///\n /// # Arguments\n /// * `llm` - The LLM to use for the execution.\n pub fn new(llm: &'a L) -> Self {\n Self { llm }\n }\n\n /// Generates a response from the LLM (non-streaming).\n ///\n /// # Arguments\n /// * `prompt` - The prompt to generate a response for.\n /// * `system_prompt` - The system prompt to use for the generation.\n ///\n /// # Returns\n /// A [Result] containing the response from the LLM or an error if there was a problem.\n pub async fn text_complete(\n &self,\n prompt: &str,\n system_prompt: &str,\n ) -> Result {\n let options = TextCompleteOptions {\n ..Default::default()\n };\n let response = self\n .llm\n .text_complete(prompt, system_prompt, options)\n .await\n .map_err(ExecutorError::Llm)?;\n Ok(ExecutorTextCompleteResponse {\n text: response.text,\n context: ExecutorContext {},\n })\n }\n\n /// Generates a streaming response from the LLM.\n ///\n /// # Arguments\n /// * `prompt` - The prompt to generate a response for.\n /// * `system_prompt` - The system prompt to use for the generation.\n ///\n /// # Returns\n /// A [Result] containing the response from the LLM or an error if there was a problem.\n pub async fn text_complete_stream(\n &self,\n prompt: &str,\n system_prompt: &str,\n ) -> Result {\n let options = TextCompleteStreamOptions {\n ..Default::default()\n };\n let response = self\n .llm\n .text_complete_stream(prompt, system_prompt, options)\n .await\n .map_err(ExecutorError::Llm)?;\n Ok(ExecutorTextCompleteStreamResponse {\n stream: response.stream,\n context: ExecutorContext {},\n })\n }\n\n /// Generates an embedding from the LLM.\n ///\n /// # Arguments\n /// * `prompt` - The item to generate an embedding for.\n ///\n /// # Returns\n ///\n /// A [Result] containing the embedding or an error if there was a problem.\n pub async fn generate_embedding(&self, prompt: &str) -> Result, ExecutorError> {\n let response = self\n .llm\n .generate_embedding(prompt)\n .await\n .map_err(ExecutorError::Llm)?;\n Ok(response)\n }\n}\n\n// TODO: Support context for completions (e.g., IDs of past conversations in Ollama).\npub struct ExecutorContext;\n\npub struct ExecutorTextCompleteResponse {\n pub text: String,\n pub context: ExecutorContext,\n}\n\npub struct ExecutorTextCompleteStreamResponse {\n pub stream: Pin> + Send>>,\n pub context: ExecutorContext,\n}\n" }, { "path": "src/lib.rs", "content": "mod core;\nmod executor;\nmod llm;\n\n// TODO: Narrow the scope of the use statements.\npub use core::*;\npub use executor::*;\npub use llm::*;\n" }, { "path": "src/llm/error.rs", "content": "use thiserror::Error;\n\nuse crate::{LlmProvider, OllamaError};\n\n#[derive(Debug, Error)]\npub enum LlmProviderError {\n #[error(\"Invalid LLM provider: {0}\")]\n InvalidValue(String),\n}\n\nimpl std::fmt::Display for LlmProvider {\n fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {\n match self {\n LlmProvider::Ollama => write!(f, \"ollama\"),\n LlmProvider::OpenAi => write!(f, \"openai\"),\n }\n }\n}\n\nimpl Default for LlmProvider {\n fn default() -> Self {\n Self::Ollama\n }\n}\n\nimpl TryFrom<&str> for LlmProvider {\n type Error = LlmProviderError;\n\n fn try_from(value: &str) -> Result {\n match value {\n \"ollama\" => Ok(LlmProvider::Ollama),\n \"openai\" => Ok(LlmProvider::OpenAi),\n _ => Err(LlmProviderError::InvalidValue(value.to_string())),\n }\n }\n}\n\n#[derive(Debug, Error)]\npub enum LlmError {\n #[error(\"Text generation error: {0}\")]\n TextGeneration(String),\n\n #[error(\"Embedding generation error: {0}\")]\n EmbeddingGeneration(String),\n\n #[error(\"Configuration error: {0}\")]\n Configuration(String),\n\n #[error(\"Ollama error: {0}\")]\n Ollama(#[from] OllamaError),\n}\n" }, { "path": "src/llm/llm_provider/mod.rs", "content": "mod ollama;\nmod openai;\n\npub use ollama::*;\n" }, { "path": "src/llm/llm_provider/ollama/config.rs", "content": "use serde::{Deserialize, Serialize};\n\n/// Default base URL for the Ollama API.\npub const DEFAULT_BASE_URL: &str = \"http://localhost:11434\";\n\n/// Default model for text completion.\npub const DEFAULT_MODEL: &str = ollama_model::;\n\n/// Default model for embeddings.\npub const DEFAULT_EMBEDDING_MODEL: &str = \"nomic-embed-text:latest\";\n\n#[derive(Debug, Clone, Serialize, Deserialize)]\npub struct OllamaConfig {\n pub base_url: Option,\n pub model: Option,\n pub embedding_model: Option,\n}\n\nimpl Default for OllamaConfig {\n fn default() -> Self {\n Self {\n base_url: Some(DEFAULT_BASE_URL.to_string()),\n model: Some(\"codestral:latest\".to_string()),\n embedding_model: Some(\"nomic-embed-text:latest\".to_string()),\n }\n }\n}\n" }, { "path": "src/llm/llm_provider/ollama/llm.rs", "content": "use thiserror::Error;\nuse tokio_stream::StreamExt;\n\nuse crate::*;\n\npub mod ollama_model {\n pub const CODESTRAL: &str = \"codestral:latest\";\n}\n\npub mod ollama_embedding_model {\n pub const NOMIC_EMBED_TEXT: &str = \"nomic-embed-text:latest\";\n}\n\n#[derive(Debug, Clone)]\npub struct Ollama<'a> {\n base_url: &'a str,\n pub model: Option<&'a str>,\n pub embeddings_model: Option<&'a str>,\n}\n\nimpl Default for Ollama<'_> {\n fn default() -> Self {\n Self {\n base_url: \"http://localhost:11434\",\n model: Some(ollama_model::CODESTRAL),\n embeddings_model: Some(ollama_embedding_model::NOMIC_EMBED_TEXT),\n }\n }\n}\n\npub struct OllamaBuilder<'a> {\n base_url: &'a str,\n model: Option<&'a str>,\n embeddings_model: Option<&'a str>,\n}\n\nimpl Default for OllamaBuilder<'_> {\n fn default() -> Self {\n let ollama = Ollama::default();\n Self {\n base_url: ollama.base_url,\n model: ollama.model,\n embeddings_model: ollama.embeddings_model,\n }\n }\n}\n\nimpl<'a> OllamaBuilder<'a> {\n pub fn new() -> Self {\n Default::default()\n }\n\n pub fn with_base_url(mut self, base_url: &'a str) -> Self {\n self.base_url = base_url;\n self\n }\n\n pub fn with_model(mut self, model: &'a str) -> Self {\n self.model = Some(model);\n self\n }\n\n pub fn with_embeddings_model(mut self, embeddings_model: &'a str) -> Self {\n self.embeddings_model = Some(embeddings_model);\n self\n }\n\n pub fn build(self) -> Ollama<'a> {\n Ollama {\n base_url: self.base_url,\n model: self.model,\n embeddings_model: self.embeddings_model,\n }\n }\n}\n\n#[derive(Error, Debug)]\npub enum OllamaError {\n #[error(\"Unexpected response from API. Error: {0}\")]\n Api(String),\n\n #[error(\"Unexpected error when parsing response from Ollama. Error: {0}\")]\n Parsing(String),\n\n #[error(\"Configuration error: {0}\")]\n Configuration(String),\n\n #[error(\"Serialization error: {0}\")]\n Serialization(String),\n\n #[error(\n \"Ollama API is not available. Please check if Ollama is running in the specified port. Error: {0}\"\n )]\n ApiUnavailable(String),\n}\n\nimpl<'a> Ollama<'a> {\n /// Lists the running models in the Ollama API.\n ///\n /// # Returns\n ///\n /// A [Result] containing the list of running models or an error if there was a problem.\n ///\n #[allow(dead_code)]\n pub(crate) fn list_running_models(&self) -> Result {\n let response = self.get_from_ollama_api(\"api/ps\")?;\n let parsed_response = Self::parse_models_response(&response)?;\n Ok(parsed_response)\n }\n\n // /// Lists the local models in the Ollama API.\n // ///\n // /// # Returns\n // ///\n // /// A [Result] containing the list of local models or an error if there was a problem.\n #[allow(dead_code)]\n pub fn list_local_models(&self) -> Result {\n let response = self.get_from_ollama_api(\"api/tags\")?;\n let parsed_response = Self::parse_models_response(&response)?;\n Ok(parsed_response)\n }\n\n fn parse_models_response(response: &str) -> Result {\n let models: OllamaApiModelsMetadata =\n serde_json::from_str(response).map_err(|e| OllamaError::Parsing(e.to_string()))?;\n Ok(models)\n }\n\n fn get_from_ollama_api(&self, url: &str) -> Result {\n let url = format!(\"{}/{}\", self.base_url()?, url);\n\n let client = reqwest::blocking::Client::new();\n let response = client\n .get(url)\n .send()\n .map_err(|e| OllamaError::ApiUnavailable(e.to_string()))?;\n let response_text = response\n .text()\n .map_err(|e| OllamaError::Api(e.to_string()))?;\n Ok(response_text)\n }\n\n fn base_url(&self) -> Result {\n Ok(self.base_url.to_string())\n }\n\n fn model(&self) -> Result {\n self.model\n .map(|s| s.to_owned())\n .ok_or_else(|| OllamaError::Configuration(\"Model not set\".to_string()))\n }\n\n fn embedding_model(&self) -> Result {\n self.embeddings_model\n .map(|s| s.to_owned())\n .ok_or_else(|| OllamaError::Configuration(\"Embedding model not set\".to_string()))\n }\n}\n\nimpl<'a> Llm for Ollama<'a> {\n async fn text_complete(\n &self,\n prompt: &str,\n system_prompt: &str,\n _options: TextCompleteOptions,\n ) -> Result {\n let body = OllamaGenerateRequest {\n model: self\n .model()\n .map_err(|_e| LlmError::Configuration(\"Model not set\".to_string()))?,\n prompt: prompt.to_string(),\n system: Some(system_prompt.to_string()),\n ..Default::default()\n };\n\n let client = reqwest::Client::new();\n let url = format!(\n \"{}/api/generate\",\n self.base_url()\n .map_err(|_e| LlmError::Configuration(\"Base URL not set\".to_string()))?\n );\n let response = client\n .post(url)\n .body(serde_json::to_string(&body).unwrap())\n .send()\n .await\n .map_err(|e| LlmError::Ollama(OllamaError::ApiUnavailable(e.to_string())))?;\n let body = response\n .text()\n .await\n .map_err(|e| LlmError::Ollama(OllamaError::Api(e.to_string())))?;\n let ollama_response: OllamaGenerateResponse = serde_json::from_str(&body)\n .map_err(|e| LlmError::Ollama(OllamaError::Parsing(e.to_string())))?;\n let response = TextCompleteResponse {\n text: ollama_response.response,\n context: ollama_response.context,\n };\n Ok(response)\n }\n\n async fn text_complete_stream(\n &self,\n prompt: &str,\n system_prompt: &str,\n options: TextCompleteStreamOptions,\n ) -> Result {\n let body = OllamaGenerateRequest {\n model: self.model()?,\n prompt: prompt.to_string(),\n stream: Some(true),\n format: None,\n images: None,\n system: Some(system_prompt.to_string()),\n keep_alive: Some(\"5m\".to_string()),\n context: options.context,\n };\n\n let url = format!(\"{}/api/generate\", self.base_url()?);\n let stream = SseClient::post(&url, Some(serde_json::to_string(&body).unwrap()));\n let stream = stream.map(|event| {\n let parsed_message = serde_json::from_str::(&event);\n match parsed_message {\n Ok(message) => Ok(message.response),\n Err(e) => Err(LlmError::Ollama(OllamaError::Parsing(e.to_string()))),\n }\n });\n let response = TextCompleteStreamResponse {\n stream: Box::pin(stream),\n };\n Ok(response)\n }\n\n async fn generate_embedding(&self, prompt: &str) -> Result, LlmError> {\n let client = reqwest::Client::new();\n let url = format!(\"{}/api/embeddings\", self.base_url()?);\n let body = OllamaEmbeddingsRequest {\n model: self.embedding_model()?,\n prompt: prompt.to_string(),\n };\n let response = client\n .post(url)\n .body(\n serde_json::to_string(&body)\n .map_err(|e| OllamaError::Serialization(e.to_string()))?,\n )\n .send()\n .await\n .map_err(|e| OllamaError::ApiUnavailable(e.to_string()))?;\n let body = response\n .text()\n .await\n .map_err(|e| OllamaError::Api(e.to_string()))?;\n let response: OllamaEmbeddingsResponse =\n serde_json::from_str(&body).map_err(|e| OllamaError::Parsing(e.to_string()))?;\n\n Ok(response.embedding)\n }\n\n fn provider(&self) -> LlmProvider {\n LlmProvider::Ollama\n }\n\n fn text_completion_model_name(&self) -> String {\n self.model().expect(\"Model not set\").to_string()\n }\n\n fn embedding_model_name(&self) -> String {\n self.embedding_model()\n .expect(\"Embedding model not set\")\n .to_string()\n }\n}\n" }, { "path": "src/llm/llm_provider/ollama/mod.rs", "content": "mod config;\nmod llm;\nmod models;\n\npub use llm::*;\npub use models::*;\n" }, { "path": "src/llm/llm_provider/ollama/models.rs", "content": "use serde::{Deserialize, Serialize};\n\nuse crate::ollama_model;\n\n/// Response from the Ollama API for obtaining information about local models.\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#list-running-models).\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaApiModelsMetadata {\n pub models: Vec,\n}\n\n/// Response item from the Ollama API for obtaining information about local models.\n///\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#response-22).\n#[allow(dead_code)]\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaApiModelMetadata {\n /// The name of the model (e.g., \"mistral:latest\")\n pub name: String,\n\n /// The Ollama identifier of the model (e.g., \"mistral:latest\")\n pub model: String,\n\n /// Size of the model in bytes\n pub size: usize,\n\n /// Digest of the model using SHA256 (e.g., \"2ae6f6dd7a3dd734790bbbf58b8909a606e0e7e97e94b7604e0aa7ae4490e6d8\")\n pub digest: String,\n\n /// Model expiry time in ISO 8601 format (e.g., \"2024-06-04T14:38:31.83753-07:00\")\n pub expires_at: Option,\n\n /// More details about the model\n pub details: OllamaApiModelDetails,\n}\n\n/// Details about a running model in the API for listing running models (`GET /api/ps`).\n///\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#response-22).\n#[allow(dead_code)]\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaApiModelDetails {\n /// Model identifier that this model is based on\n pub parent_model: String,\n\n /// Format that this model is stored in (e.g., \"gguf\")\n pub format: String,\n\n /// Model family (e.g., \"ollama\")\n pub family: String,\n\n /// Parameters of the model (e.g., \"7.2B\")\n pub parameter_size: String,\n\n /// Quantization level of the model (e.g., \"Q4_0\" for 4-bit quantization)\n pub quantization_level: String,\n}\n\n/// Request for generating a response from the Ollama API.\n///\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#generate-a-completion).\n#[allow(dead_code)]\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaGenerateRequest {\n /// Model identifier (e.g., \"mistral:latest\")\n pub model: String,\n\n /// The prompt to generate a response for (e.g., \"List all Kubernetes pods\")\n pub prompt: String,\n\n /// The context parameter returned from a previous request to /generate, this can be used to keep a short conversational memory\n pub context: Option>,\n\n /// Optional list of base64-encoded images (for multimodal models such as `llava`)\n pub images: Option>,\n\n /// Optional format to use for the response (currently only \"json\" is supported)\n pub format: Option,\n\n /// Optional flag that controls whether the response is streamed or not (defaults to true).\n /// If `false`` the response will be returned as a single response object, rather than a stream of objects\n pub stream: Option,\n\n // System message (overrides what is defined in the Modelfile)\n pub system: Option,\n\n /// Controls how long the model will stay loaded into memory following the request (default: 5m)\n pub keep_alive: Option,\n}\n\nimpl Default for OllamaGenerateRequest {\n fn default() -> Self {\n Self {\n model: ollama_model::CODESTRAL.to_string(),\n prompt: \"\".to_string(),\n stream: Some(false),\n format: None,\n images: None,\n system: Some(\"You are a helpful assistant\".to_string()),\n keep_alive: Some(\"5m\".to_string()),\n context: None,\n }\n }\n}\n\n#[derive(Debug, Serialize, Deserialize)]\n#[allow(dead_code)]\npub struct OllamaGenerateResponse {\n /// Model identifier (e.g., \"mistral:latest\")\n pub model: String,\n\n /// Time at which the response was generated (ISO 8601 format)\n pub created_at: String,\n\n /// The response to the prompt\n pub response: String,\n\n /// The encoding of the conversation used in this response, this can be sent in the next request to keep a conversational memory\n pub context: Option>,\n\n /// The duration of the response in nanoseconds\n pub total_duration: usize,\n}\n\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaGenerateStreamItemResponse {\n /// Model identifier (e.g., \"mistral:latest\")\n pub model: String,\n\n /// Time at which the response was generated (ISO 8601 format)\n pub created_at: String,\n\n /// The response to the prompt\n pub response: String,\n}\n\n/// Request for generating an embedding from the Ollama API.\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#generate-embeddings).\n///\n#[allow(dead_code)]\n#[derive(Debug, Serialize, Deserialize, Default)]\npub struct OllamaEmbeddingsRequest {\n /// The string to generate an embedding for.\n pub prompt: String,\n\n /// The model to use for the embedding generation.\n pub model: String,\n}\n\n/// Response from the Ollama API for generating an embedding.\n/// Referenced from the Ollama API documentation [here](https://github.com/ollama/ollama/blob/fedf71635ec77644f8477a86c6155217d9213a11/docs/api.md#generate-embeddings).\n///\n#[allow(dead_code)]\n#[derive(Debug, Serialize, Deserialize)]\npub struct OllamaEmbeddingsResponse {\n /// The embedding for the prompt.\n pub embedding: Vec,\n}\n" }, { "path": "src/llm/llm_provider/openai.rs", "content": "// use async_trait::async_trait;\n// use openai_api_rs::v1::{\n// api::OpenAIClient,\n// chat_completion::{self, ChatCompletionRequest},\n// common::{GPT3_5_TURBO, GPT4, GPT4_O},\n// };\n\n// pub mod openai_model {\n// pub const GPT35_TURBO: &str = GPT35_TURBO;\n// pub const GPT4: &str = GPT4;\n// pub const GPT40: &str = GPT40;\n// }\n\n// pub struct OpenAi<'a> {\n// pub model: &'a str,\n// api_key: &'a str,\n// }\n\n// impl<'a> OpenAi<'a> {\n// pub fn new(api_key: &'a str, model: &'a str) -> Self {\n// Self { api_key, model }\n// }\n// }\n\n// #[async_trait]\n// impl<'a> TextCompletionLlm for OpenAi<'a> {\n// async fn complete(\n// &self,\n// system_prompts: &[String],\n// ) -> Result> {\n// let client = OpenAIClient::new(self.api_key.to_owned());\n// let system_msgs = system_prompts\n// .iter()\n// .map(|p| chat_completion::ChatCompletionMessage {\n// role: chat_completion::MessageRole::system,\n// content: chat_completion::Content::Text(p.to_owned()),\n// name: None,\n// tool_calls: None,\n// tool_call_id: None,\n// })\n// .collect::>();\n// let mut req = ChatCompletionRequest::new(self.model.to_owned(), system_msgs);\n// req.max_tokens = Some(self.config.max_tokens as i64);\n// req.temperature = Some(self.config.temperature);\n\n// let result = client.chat_completion(req).await?;\n// let completion = result\n// .choices\n// .first()\n// .unwrap()\n// .message\n// .content\n// .clone()\n// .unwrap();\n// Ok(completion)\n// }\n// }\n" }, { "path": "src/llm/mod.rs", "content": "mod error;\nmod llm_provider;\nmod models;\n\npub use error::*;\npub use llm_provider::*;\npub use models::*;\n" }, { "path": "src/llm/models.rs", "content": "#![allow(dead_code)]\n\nuse std::pin::Pin;\n\nuse dyn_clone::DynClone;\nuse serde::{Deserialize, Serialize};\nuse tokio_stream::Stream;\n\nuse super::error::LlmError;\n\n#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]\npub enum LlmProvider {\n #[serde(rename = \"ollama\")]\n Ollama,\n #[serde(rename = \"openai\")]\n OpenAi,\n}\n\n/// A trait for LLM providers which implements text completion, embeddings, etc.\n///\n/// > `DynClone` is used so that there can be dynamic dispatch of the `Llm` trait,\n/// > especially needed for [magic-cli](https://github.com/guywaldman/magic-cli).\npub trait Llm: DynClone {\n /// Generates a response from the LLM.\n ///\n /// # Arguments\n /// * `prompt` - The prompt to generate a response for.\n /// * `system_prompt` - The system prompt to use for the generation.\n /// * `options` - The options for the generation.\n ///\n /// # Returns\n /// A [Result] containing the response from the LLM or an error if there was a problem.\n ///\n fn text_complete(\n &self,\n prompt: &str,\n system_prompt: &str,\n options: TextCompleteOptions,\n ) -> impl std::future::Future> + Send;\n\n /// Generates a streaming response from the LLM.\n ///\n /// # Arguments\n /// * `prompt` - The prompt to generate a response for.\n /// * `system_prompt` - The system prompt to use for the generation.\n /// * `options` - The options for the generation.\n ///\n /// # Returns\n /// A [Result] containing the response from the LLM or an error if there was a problem.\n ///\n fn text_complete_stream(\n &self,\n prompt: &str,\n system_prompt: &str,\n options: TextCompleteStreamOptions,\n ) -> impl std::future::Future> + Send;\n\n /// Generates an embedding from the LLM.\n ///\n /// # Arguments\n /// * `prompt` - The item to generate an embedding for.\n ///\n /// # Returns\n ///\n /// A [Result] containing the embedding or an error if there was a problem.\n fn generate_embedding(\n &self,\n prompt: &str,\n ) -> impl std::future::Future, LlmError>> + Send;\n\n /// Returns the provider of the LLM.\n fn provider(&self) -> LlmProvider;\n\n /// Returns the name of the model used for text completions.\n fn text_completion_model_name(&self) -> String;\n\n /// Returns the name of the model used for embeddings.\n fn embedding_model_name(&self) -> String;\n}\n\n#[derive(Debug, Clone, Default)]\npub struct TextCompleteOptions {\n /// An encoding of the conversation used in this response, this can be sent in the next request to keep a conversational memory.\n /// This should be as returned from the previous response.\n pub context: Option>,\n}\n\n#[derive(Debug, Clone, Default)]\npub struct TextCompleteStreamOptions {\n pub context: Option>,\n}\n\n#[derive(Debug, Clone)]\npub struct TextCompleteResponse {\n pub text: String,\n // TODO: This is specific to Ollama, context looks differently for other LLM providers.\n pub context: Option>,\n}\n\npub struct TextCompleteStreamResponse {\n pub stream: Pin> + Send>>,\n // TODO: Handle context with streaming response.\n // pub context: Vec,\n}\n\n#[derive(Debug)]\npub(crate) struct SystemPromptResponseOption {\n pub scenario: String,\n pub type_name: String,\n pub response: String,\n pub schema: Vec,\n}\n\n#[derive(Debug)]\npub(crate) struct SystemPromptCommandSchemaField {\n pub name: String,\n pub description: String,\n pub typ: String,\n pub example: String,\n}\n" } ]