[ { "path": ".github/ISSUE_TEMPLATE/bug_report.md", "content": "---\nname: Bug report\nabout: Create a report to help us improve\ntitle: ''\nlabels: ''\nassignees: ''\n\n---\n\n\n\n**Describe the bug**\n\n\n**To Reproduce**\n\n\n**Expected behavior**\n\n\n**Logs**\n\n\n**Screenshots**\n\n\n**Configuration**\n\n\n**Environment (please complete the following information):**\n- os version: [e.g. Ubuntu 20.04]\n- aichat version: [e.g. 0.9.0]\n- terminal version: [e.g. GNOME Terminal 3.44.0]\n\n**Additional context**\n" }, { "path": ".github/ISSUE_TEMPLATE/feature_request.md", "content": "---\nname: Feature request\nabout: Suggest an idea for this project\ntitle: ''\nlabels: enhancement\nassignees: ''\n\n---\n\n\n\n**Is your feature request related to a problem? Please describe.**\n\n\n**Describe the solution you'd like**\n\n\n**Describe alternatives you've considered**\n\n\n**Additional context**\n" }, { "path": ".github/workflows/ci.yaml", "content": "name: CI\n\non:\n pull_request:\n branches:\n - '*'\n push:\n branches:\n - main\n\ndefaults:\n run:\n shell: bash\n\njobs:\n all:\n name: All\n\n strategy:\n matrix:\n os:\n - ubuntu-latest\n - macos-latest\n - windows-latest\n\n runs-on: ${{matrix.os}}\n\n env:\n RUSTFLAGS: --deny warnings\n\n steps:\n - uses: actions/checkout@v4\n\n - name: Install Rust Toolchain Components\n uses: dtolnay/rust-toolchain@stable\n\n - uses: Swatinem/rust-cache@v2\n\n - name: Test\n run: cargo test --all\n\n - name: Clippy\n run: cargo clippy --all --all-targets -- -D warnings\n\n - name: Format\n run: cargo fmt --all --check" }, { "path": ".github/workflows/release.yaml", "content": "name: Release\n\non:\n push:\n tags:\n - v[0-9]+.[0-9]+.[0-9]+*\n\njobs:\n release:\n name: Publish to GitHub Release\n permissions:\n contents: write\n outputs:\n rc: ${{ steps.check-tag.outputs.rc }}\n\n strategy:\n matrix:\n include:\n - target: aarch64-unknown-linux-musl\n os: ubuntu-latest\n use-cross: true\n cargo-flags: \"\"\n - target: aarch64-apple-darwin\n os: macos-latest\n use-cross: true\n cargo-flags: \"\"\n - target: aarch64-pc-windows-msvc\n os: windows-latest\n use-cross: true\n cargo-flags: \"\"\n - target: x86_64-apple-darwin\n os: macos-latest\n cargo-flags: \"\"\n - target: x86_64-pc-windows-msvc\n os: windows-latest\n cargo-flags: \"\"\n - target: x86_64-unknown-linux-musl\n os: ubuntu-latest\n use-cross: true\n cargo-flags: \"\"\n - target: i686-unknown-linux-musl\n os: ubuntu-latest\n use-cross: true\n cargo-flags: \"\"\n - target: i686-pc-windows-msvc\n os: windows-latest\n use-cross: true\n cargo-flags: \"\"\n - target: armv7-unknown-linux-musleabihf\n os: ubuntu-latest\n use-cross: true\n cargo-flags: \"\"\n - target: arm-unknown-linux-musleabihf\n os: ubuntu-latest\n use-cross: true\n cargo-flags: \"\"\n\n runs-on: ${{matrix.os}}\n env:\n BUILD_CMD: cargo\n\n steps:\n - uses: actions/checkout@v4\n\n - name: Check Tag\n id: check-tag\n shell: bash\n run: |\n ver=${GITHUB_REF##*/}\n echo \"version=$ver\" >> $GITHUB_OUTPUT\n if [[ \"$ver\" =~ [0-9]+.[0-9]+.[0-9]+$ ]]; then\n echo \"rc=false\" >> $GITHUB_OUTPUT\n else\n echo \"rc=true\" >> $GITHUB_OUTPUT\n fi\n\n\n - name: Install Rust Toolchain Components\n uses: dtolnay/rust-toolchain@stable\n with:\n targets: ${{ matrix.target }}\n\n - name: Install cross\n if: matrix.use-cross\n uses: taiki-e/install-action@v2\n with:\n tool: cross\n\n - name: Overwrite build command env variable\n if: matrix.use-cross\n shell: bash\n run: echo \"BUILD_CMD=cross\" >> $GITHUB_ENV\n \n - name: Show Version Information (Rust, cargo, GCC)\n shell: bash\n run: |\n gcc --version || true\n rustup -V\n rustup toolchain list\n rustup default\n cargo -V\n rustc -V\n \n - name: Build\n shell: bash\n run: $BUILD_CMD build --locked --release --target=${{ matrix.target }} ${{ matrix.cargo-flags }}\n\n - name: Build Archive\n shell: bash\n id: package\n env:\n target: ${{ matrix.target }}\n version: ${{ steps.check-tag.outputs.version }}\n run: |\n set -euxo pipefail\n\n bin=${GITHUB_REPOSITORY##*/}\n dist_dir=`pwd`/dist\n name=$bin-$version-$target\n executable=target/$target/release/$bin\n\n if [[ \"$RUNNER_OS\" == \"Windows\" ]]; then\n executable=$executable.exe\n fi\n\n mkdir $dist_dir\n cp $executable $dist_dir\n cd $dist_dir\n\n if [[ \"$RUNNER_OS\" == \"Windows\" ]]; then\n archive=$dist_dir/$name.zip\n 7z a $archive *\n echo \"archive=dist/$name.zip\" >> $GITHUB_OUTPUT\n else\n archive=$dist_dir/$name.tar.gz\n tar -czf $archive *\n echo \"archive=dist/$name.tar.gz\" >> $GITHUB_OUTPUT\n fi\n\n - name: Publish Archive\n uses: softprops/action-gh-release@v2\n if: ${{ startsWith(github.ref, 'refs/tags/') }}\n with:\n draft: false\n files: ${{ steps.package.outputs.archive }}\n prerelease: ${{ steps.check-tag.outputs.rc == 'true' }}\n\n publish-crate:\n name: Publish to crates.io\n if: ${{ needs.release.outputs.rc == 'false' }}\n runs-on: ubuntu-latest\n needs: release\n steps:\n - uses: actions/checkout@v4\n\n - uses: dtolnay/rust-toolchain@stable\n\n - name: Publish\n env:\n CARGO_REGISTRY_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}\n run: cargo publish" }, { "path": ".gitignore", "content": "/target\n/tmp\n/.env\n*.log" }, { "path": "Argcfile.sh", "content": "#!/usr/bin/env bash\nset -e\n\n# @meta dotenv\n# @env DRY_RUN Dry run mode\n\n# @cmd Test configuration initialization\n# @env AICHAT_CONFIG_DIR=tmp/test-init-config\n# @arg args~\ntest-init-config() {\n unset OPENAI_API_KEY\n mkdir -p \"$AICHAT_CONFIG_DIR\"\n config_file=\"$AICHAT_CONFIG_DIR/config.yaml\"\n if [[ -f \"$config_file\" ]]; then\n rm -f \"$config_file\"\n fi\n cargo run -- \"$@\"\n}\n\n# @cmd Test running without configuration file\n# @env AICHAT_PROVIDER!\n# @env AICHAT_CONFIG_DIR=tmp/test-provider-env\n# @arg args~\ntest-no-config() {\n mkdir -p \"$AICHAT_CONFIG_DIR\"\n rm -rf \"$AICHAT_CONFIG_DIR/config.yaml\"\n cargo run -- \"$@\"\n}\n\n# @cmd Test function calling\n# @option -m --model[?`_choice_model`]\n# @option -p --preset[=weather|multi-weathers]\n# @flag -S --no-stream\n# @arg text~\ntest-function-calling() {\n args=(--role %functions%)\n if [[ -n \"$argc_model\" ]]; then\n args+=(\"--model\" \"$argc_model\")\n fi\n if [[ -n \"$argc_no_stream\" ]]; then\n args+=(\"-S\")\n fi\n if [[ -z \"$argc_text\" ]]; then\n case \"$argc_preset\" in\n multi-weathers)\n text=\"what is the weather in London and Pairs?\"\n ;;\n weather|*)\n text=\"what is the weather in London?\"\n ;;\n esac\n else\n text=\"${argc_text[*]}\"\n fi\n cargo run -- \"${args[@]}\" \"$text\"\n}\n\n# @cmd Test clients\n# @arg clients+[`_choice_client`]\ntest-clients() {\n for c in \"${argc_clients[@]}\"; do\n echo \"### $c stream\"\n aichat -m \"$c\" 1 + 2 = ?\n echo \"### $c non-stream\"\n aichat -m \"$c\" -S 1 + 2 = ?\n done\n}\n\n# @cmd Test proxy server\n# @option -m --model[?`_choice_model`]\n# @flag -S --no-stream\n# @arg text~\ntest-server() {\n args=()\n if [[ -n \"$argc_no_stream\" ]]; then\n args+=(\"-S\")\n fi\n argc chat-llm \"${args[@]}\" \\\n --api-base http://localhost:8000/v1 \\\n --model \"${argc_model:-default}\" \\\n \"$@\"\n}\n\n# @cmd Chat with any LLM api \n# @flag -S --no-stream\n# @arg provider_model![?`_choice_provider_model`]\n# @arg text~\nchat() {\n if [[ \"$argc_provider_model\" == *':'* ]]; then\n model=\"${argc_provider_model##*:}\"\n argc_provider=\"${argc_provider_model%:*}\"\n else\n argc_provider=\"${argc_provider_model}\"\n fi\n for provider_config in \"${OPENAI_COMPATIBLE_PROVIDERS[@]}\"; do\n if [[ \"$argc_provider\" == \"${provider_config%%,*}\" ]]; then\n _retrieve_api_base\n break\n fi\n done\n if [[ -n \"$api_base\" ]]; then\n env_prefix=\"$(echo \"$argc_provider\" | tr '[:lower:]' '[:upper:]')\"\n api_key_env=\"${env_prefix}_API_KEY\"\n api_key=\"${!api_key_env}\" \n if [[ -z \"$model\" ]]; then\n model=\"$(echo \"$provider_config\" | cut -d, -f2)\"\n fi\n if [[ -z \"$model\" ]]; then\n model_env=\"${env_prefix}_MODEL\"\n model=\"${!model_env}\"\n fi\n argc chat-openai-compatible \\\n --api-base \"$api_base\" \\\n --api-key \"$api_key\" \\\n --model \"$model\" \\\n \"${argc_text[@]}\"\n else\n argc chat-$argc_provider \"${argc_text[@]}\"\n fi\n}\n\n# @cmd List models by openai-compatible api\n# @flag --name-only Print model name only\n# @arg provider![`_choice_provider`]\nmodels() {\n for provider_config in \"${OPENAI_COMPATIBLE_PROVIDERS[@]}\"; do\n if [[ \"$argc_provider\" == \"${provider_config%%,*}\" ]]; then\n _retrieve_api_base\n break\n fi\n done\n if [[ -n \"$api_base\" ]]; then\n env_prefix=\"$(echo \"$argc_provider\" | tr '[:lower:]' '[:upper:]')\"\n api_key_env=\"${env_prefix}_API_KEY\"\n api_key=\"${!api_key_env}\" \n jq_args=()\n if [[ -n \"$argc_name_only\" ]]; then\n case \"$argc_provider\" in\n cloudflare)\n jq_args+=(-r '.result[].name')\n ;;\n github)\n jq_args+=(-r '.[].name')\n ;;\n *)\n jq_args+=(-r '.data[].id')\n ;;\n esac\n fi\n _openai_compatible_models | jq \"${jq_args[@]}\"\n else\n if ! cat \"$0\" | grep -q \"^models-$argc_provider\"; then\n _die \"error: provider '$argc_provider' does not have a models api\"\n fi\n cli_args=()\n if [[ -n \"$argc_name_only\" ]]; then\n cli_args+=(--name-only)\n fi\n argc models-$argc_provider \"${cli_args[@]}\"\n fi\n}\n\n# @cmd Chat with openai-compatible api\n# @option --api-base! $$ \n# @option --api-key! $$\n# @option -m --model! $$\n# @flag -S --no-stream\n# @arg text~\nchat-openai-compatible() {\n _wrapper curl -i \"$argc_api_base/chat/completions\" \\\n-X POST \\\n-H \"Content-Type: application/json\" \\\n-H \"Authorization: Bearer $argc_api_key\" \\\n-d \"$(_build_body openai \"$@\")\"\n}\n\n# @cmd List models by openai-compatible api\n# @option --api-base! $$\n# @option --api-key! $$\n# @flag --name-only Print model name only\nmodels-openai-compatible() {\n jq_args=()\n if [[ -n \"$argc_name_only\" ]]; then\n jq_args+=(-r '.data[].id')\n fi\n _openai_compatible_models | jq \"${jq_args[@]}\"\n}\n\n# @cmd Chat with azure-openai api\n# @option --api-url! $$ \n# @option --api-key! $$\n# @option -m --model! $$\n# @flag -S --no-stream\n# @arg text~\nchat-azure-openai() {\n _wrapper curl -i \"$argc_api_url\" \\\n-X POST \\\n-H \"Content-Type: application/json\" \\\n-H \"api-key: $argc_api_key\" \\\n-d \"$(_build_body openai \"$@\")\"\n}\n\n# @cmd Chat with gemini api\n# @env GEMINI_API_KEY!\n# @option -m --model=gemini-1.5-pro-latest $GEMINI_MODEL\n# @flag -S --no-stream\n# @arg text~\nchat-gemini() {\n method=\"streamGenerateContent\"\n if [[ -n \"$argc_no_stream\" ]]; then\n method=\"generateContent\"\n fi\n _wrapper curl -i \"https://generativelanguage.googleapis.com/v1beta/models/${argc_model}:${method}?key=${GEMINI_API_KEY}\" \\\n-i -X POST \\\n-H 'Content-Type: application/json' \\\n-d \"$(_build_body gemini \"$@\")\" \n}\n\n# @cmd List gemini models\n# @env GEMINI_API_KEY!\n# @flag --name-only Print model name only\nmodels-gemini() {\n jq_args=()\n if [[ -n \"$argc_name_only\" ]]; then\n jq_args+=(-r '.models[].name')\n fi\n _wrapper curl -fsSL \"https://generativelanguage.googleapis.com/v1beta/models?key=${GEMINI_API_KEY}\" \\\n-H 'Content-Type: application/json' \\\n | jq \"${jq_args[@]}\"\n}\n\n# @cmd Chat with claude api\n# @env CLAUDE_API_KEY!\n# @option -m --model=claude-3-haiku-20240307 $CLAUDE_MODEL\n# @flag -S --no-stream\n# @arg text~\nchat-claude() {\n _wrapper curl -i https://api.anthropic.com/v1/messages \\\n-X POST \\\n-H 'content-type: application/json' \\\n-H 'anthropic-version: 2023-06-01' \\\n-H 'anthropic-beta: tools-2024-05-16' \\\n-H \"x-api-key: $CLAUDE_API_KEY\" \\\n-d \"$(_build_body claude \"$@\")\"\n}\n\n# @cmd List claude models\n# @env CLAUDE_API_KEY!\n# @flag --name-only Print model name only\nmodels-claude() {\n jq_args=()\n if [[ -n \"$argc_name_only\" ]]; then\n jq_args+=(-r '.data[].id')\n fi\n _wrapper curl -fsSL \"https://api.anthropic.com/v1/models\" \\\n-H 'Content-Type: application/json' \\\n-H 'anthropic-version: 2023-06-01' \\\n-H \"x-api-key: $CLAUDE_API_KEY\" \\\n | jq \"${jq_args[@]}\"\n}\n\n# @cmd Chat with cohere api\n# @env COHERE_API_KEY!\n# @option -m --model=command-r-08-2024 $COHERE_MODEL\n# @flag -S --no-stream\n# @arg text~\nchat-cohere() {\n _wrapper curl -i https://api.cohere.ai/v2/chat \\\n-X POST \\\n-H 'Content-Type: application/json' \\\n-H \"Authorization: Bearer $COHERE_API_KEY\" \\\n-d \"$(_build_body cohere \"$@\")\"\n}\n\n# @cmd List cohere models\n# @env COHERE_API_KEY!\n# @flag --name-only Print model name only\nmodels-cohere() {\n jq_args=()\n if [[ -n \"$argc_name_only\" ]]; then\n jq_args+=(-r '.models[].name')\n fi\n _wrapper curl -fsSL https://api.cohere.ai/v1/models \\\n-H \"Authorization: Bearer $COHERE_API_KEY\" \\\n | jq \"${jq_args[@]}\"\n}\n\n# @cmd Chat with vertexai api\n# @env require-tools gcloud\n# @env VERTEXAI_PROJECT_ID!\n# @env VERTEXAI_LOCATION!\n# @option -m --model=gemini-1.5-flash-002 $VERTEXAI_GEMINI_MODEL\n# @flag -S --no-stream\n# @arg text~\nchat-vertexai() {\n api_key=\"$(gcloud auth print-access-token)\"\n func=\"streamGenerateContent\"\n if [[ -n \"$argc_no_stream\" ]]; then\n func=\"generateContent\"\n fi\n url=https://$VERTEXAI_LOCATION-aiplatform.googleapis.com/v1/projects/$VERTEXAI_PROJECT_ID/locations/$VERTEXAI_LOCATION/publishers/google/models/$argc_model:$func\n _wrapper curl -i $url \\\n-X POST \\\n-H \"Authorization: Bearer $api_key\" \\\n-H 'Content-Type: application/json' \\\n-d \"$(_build_body vertexai \"$@\")\" \n}\n\n_argc_before() {\n OPENAI_COMPATIBLE_PROVIDERS=( \\\n openai,gpt-4o-mini,https://api.openai.com/v1 \\\n ai21,jamba-1.5-mini,https://api.ai21.com/studio/v1 \\\n cloudflare,@cf/meta/llama-3.1-8b-instruct,https://api.cloudflare.com/client/v4/accounts/${CLOUDFLARE_ACCOUNT_ID}/ai/v1 \\\n deepinfra,meta-llama/Meta-Llama-3.1-8B-Instruct,https://api.deepinfra.com/v1/openai \\\n deepseek,deepseek-chat,https://api.deepseek.com \\\n ernie,ernie-4.0-turbo-8k-latest,https://qianfan.baidubce.com/v2 \\\n github,gpt-4o-mini,https://models.inference.ai.azure.com \\\n groq,llama-3.1-8b-instant,https://api.groq.com/openai/v1 \\\n hunyuan,hunyuan-large,https://api.hunyuan.cloud.tencent.com/v1 \\\n minimax,MiniMax-Text-01,https://api.minimax.chat/v1 \\\n mistral,mistral-small-latest,https://api.mistral.ai/v1 \\\n moonshot,moonshot-v1-8k,https://api.moonshot.cn/v1 \\\n openrouter,openai/gpt-4o-mini,https://openrouter.ai/api/v1 \\\n perplexity,llama-3.1-8b-instruct,https://api.perplexity.ai \\\n qianwen,qwen-turbo-latest,https://dashscope.aliyuncs.com/compatible-mode/v1 \\\n xai,grok-beta,https://api.x.ai/v1 \\\n zhipuai,glm-4-0520,https://open.bigmodel.cn/api/paas/v4 \\\n )\n\n stream=\"true\"\n if [[ -n \"$argc_no_stream\" ]]; then\n stream=\"false\"\n fi\n}\n\n_openai_compatible_models() {\n api_base=\"${api_base:-\"$argc_api_base\"}\"\n api_key=\"${api_key:-\"$argc_api_key\"}\"\n url=\"${api_base}/models\"\n if [[ \"$argc_provider\" == \"cloudflare\" ]]; then\n url=\"https://api.cloudflare.com/client/v4/accounts/${CLOUDFLARE_ACCOUNT_ID}/ai/models/search\"\n fi\n\n _wrapper curl -fsSL \"$url\" \\\n-H \"Authorization: Bearer $api_key\" \\\n\n}\n\n_retrieve_api_base() {\n api_base=\"${provider_config##*,}\"\n if [[ -z \"$api_base\" ]]; then\n key=\"$(echo $argc_provider | tr '[:lower:]' '[:upper:]')_API_BASE\"\n api_base=\"${!key}\"\n if [[ -z \"$api_base\" ]]; then\n _die \"error: miss api_base for $argc_provider; please set $key\"\n fi\n fi\n}\n\n_choice_model() {\n aichat --list-models\n}\n\n_choice_provider_model() {\n _choice_provider\n _choice_model\n}\n\n_choice_provider() {\n _choice_client\n _choice_openai_compatible_provider\n}\n\n_choice_client() {\n printf \"%s\\n\" gemini claude cohere azure-openai vertexai bedrock\n}\n\n_choice_openai_compatible_provider() {\n for provider_config in \"${OPENAI_COMPATIBLE_PROVIDERS[@]}\"; do\n echo \"${provider_config%%,*}\"\n done\n}\n\n_build_body() {\n kind=\"$1\"\n if [[ \"$#\" -eq 1 ]]; then\n file=\"${BODY_FILE:-\"tmp/body/$1.json\"}\"\n if [[ -f \"$file\" ]]; then\n cat \"$file\" | \\\n sed -E \\\n -e 's%\"model\": \".*\"%\"model\": \"'\"$argc_model\"'\"%' \\\n -e 's%\"stream\": (true|false)%\"stream\": '$stream'%' \\\n\n fi\n else\n shift\n case \"$kind\" in\n openai|cohere)\n echo '{\n \"model\": \"'$argc_model'\",\n \"messages\": [\n {\n \"role\": \"user\",\n \"content\": \"'\"$*\"'\"\n }\n ],\n \"stream\": '$stream'\n}'\n ;;\n claude)\n echo '{\n \"model\": \"'$argc_model'\",\n \"messages\": [\n {\n \"role\": \"user\",\n \"content\": \"'\"$*\"'\"\n }\n ],\n \"max_tokens\": 4096,\n \"stream\": '$stream'\n}'\n\n ;;\n gemini|vertexai)\n echo '{\n \"contents\": [{\n \"role\": \"user\",\n \"parts\": [\n {\n \"text\": \"'\"$*\"'\"\n }\n ]\n }],\n \"safetySettings\":[{\"category\":\"HARM_CATEGORY_HARASSMENT\",\"threshold\":\"BLOCK_ONLY_HIGH\"},{\"category\":\"HARM_CATEGORY_HATE_SPEECH\",\"threshold\":\"BLOCK_ONLY_HIGH\"},{\"category\":\"HARM_CATEGORY_SEXUALLY_EXPLICIT\",\"threshold\":\"BLOCK_ONLY_HIGH\"},{\"category\":\"HARM_CATEGORY_DANGEROUS_CONTENT\",\"threshold\":\"BLOCK_ONLY_HIGH\"}]\n}'\n ;;\n *)\n _die \"error: unsupported build body for $kind\"\n ;;\n esac\n\n fi\n}\n\n_wrapper() {\n if [[ \"$DRY_RUN\" == \"true\" ]] || [[ \"$DRY_RUN\" == \"1\" ]]; then\n echo \"$@\" >&2\n else\n \"$@\"\n fi\n}\n\n_die() {\n echo $*\n exit 1\n}\n\n# See more details at https://github.com/sigoden/argc\neval \"$(argc --argc-eval \"$0\" \"$@\")\"\n" }, { "path": "Cargo.toml", "content": "[package]\nname = \"aichat\"\nversion = \"0.30.0\"\nedition = \"2021\"\nauthors = [\"sigoden \"]\ndescription = \"All-in-one LLM CLI Tool\"\nlicense = \"MIT OR Apache-2.0\"\nhomepage = \"https://github.com/sigoden/aichat\"\nrepository = \"https://github.com/sigoden/aichat\"\ncategories = [\"command-line-utilities\"]\nkeywords = [\"chatgpt\", \"llm\", \"cli\", \"ai\", \"repl\"]\n\n[dependencies]\nanyhow = \"1.0.69\"\nbytes = \"1.4.0\"\nclap = { version = \"4.4.8\", features = [\"derive\"] }\ndirs = \"6.0.0\"\nfutures-util = \"0.3.29\"\ninquire = \"0.7.0\"\nis-terminal = \"0.4.9\"\nreedline = \"0.40.0\"\nserde = { version = \"1.0.152\", features = [\"derive\"] }\nserde_json = { version = \"1.0.93\", features = [\"preserve_order\"] }\nserde_yaml = \"0.9.17\"\ntokio = { version = \"1.34.0\", features = [\"rt\", \"time\", \"macros\", \"signal\", \"rt-multi-thread\"] }\ntokio-graceful = \"0.2.2\"\ntokio-stream = { version = \"0.1.15\", default-features = false, features = [\"sync\"] }\ncrossterm = \"0.28.1\"\nchrono = \"0.4.23\"\nbincode = { version = \"2.0.0\", features = [\"serde\", \"std\"], default-features = false }\nparking_lot = \"0.12.1\"\nfancy-regex = \"0.14.0\"\nbase64 = \"0.22.0\"\nnu-ansi-term = \"0.50.0\"\nasync-trait = \"0.1.74\"\ntextwrap = \"0.16.0\"\nansi_colours = \"1.2.2\"\nreqwest-eventsource = \"0.6.0\"\nsimplelog = \"0.12.1\"\nlog = \"0.4.20\"\nshell-words = \"1.1.0\"\nsha2 = \"0.10.8\"\nunicode-width = \"0.2.0\"\nasync-recursion = \"1.1.1\"\nhttp = \"1.1.0\"\nhttp-body-util = \"0.1\"\nhyper = { version = \"1.0\", features = [\"full\"] }\nhyper-util = { version = \"0.1\", features = [\"server-auto\", \"client-legacy\"] }\ntime = { version = \"0.3.36\", features = [\"macros\"] }\nindexmap = { version = \"2.2.6\", features = [\"serde\"] }\nhmac = \"0.12.1\"\naws-smithy-eventstream = \"0.60.4\"\nurlencoding = \"2.1.3\"\nunicode-segmentation = \"1.11.0\"\njson-patch = { version = \"4.0.0\", default-features = false }\nbitflags = \"2.5.0\"\npath-absolutize = \"3.1.1\"\nhnsw_rs = \"0.3.0\"\nrayon = \"1.10.0\"\nuuid = { version = \"1.9.1\", features = [\"v4\"] }\nscraper = { version = \"0.23.1\", default-features = false, features = [\"deterministic\"] }\nsys-locale = \"0.3.1\"\nhtml_to_markdown = \"0.1.0\"\nrust-embed = \"8.5.0\"\nos_info = { version = \"3.8.2\", default-features = false }\nbm25 = { version = \"2.0.1\", features = [\"parallelism\"] }\nwhich = \"8.0.0\"\nfuzzy-matcher = \"0.3.7\"\nterminal-colorsaurus = \"0.4.8\"\nduct = \"1.0.0\"\n\n[dependencies.reqwest]\nversion = \"0.12.0\"\nfeatures = [\"json\", \"multipart\", \"socks\", \"rustls-tls\", \"rustls-tls-native-roots\"]\ndefault-features = false\n\n[dependencies.syntect]\nversion = \"5.0.0\"\ndefault-features = false\nfeatures = [\"parsing\", \"regex-onig\", \"plist-load\"]\n\n[target.'cfg(target_os = \"macos\")'.dependencies]\ncrossterm = { version = \"0.28.1\", features = [\"use-dev-tty\"] }\n\n[target.'cfg(target_os = \"linux\")'.dependencies]\narboard = { version = \"3.3.0\", default-features = false, features = [\"wayland-data-control\"] }\n\n[target.'cfg(not(any(target_os = \"linux\", target_os = \"android\", target_os = \"emscripten\")))'.dependencies]\narboard = { version = \"3.3.0\", default-features = false }\n\n[dev-dependencies]\npretty_assertions = \"1.4.0\"\nrand = \"0.9.0\"\n\n[profile.release]\nlto = true\nstrip = true\nopt-level = \"z\"\n" }, { "path": "LICENSE-APACHE", "content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n" }, { "path": "LICENSE-MIT", "content": "The MIT License (MIT)\n\nCopyright (c) sigoden\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": "# AIChat: All-in-one LLM CLI Tool\n\n[![CI](https://github.com/sigoden/aichat/actions/workflows/ci.yaml/badge.svg)](https://github.com/sigoden/aichat/actions/workflows/ci.yaml)\n[![Crates](https://img.shields.io/crates/v/aichat.svg)](https://crates.io/crates/aichat)\n[![Discord](https://img.shields.io/discord/1226737085453701222?label=Discord)](https://discord.gg/mr3ZZUB9hG)\n\nAIChat is an all-in-one LLM CLI tool featuring Shell Assistant, CMD & REPL Mode, RAG, AI Tools & Agents, and More. \n\n## Install\n\n### Package Managers\n\n- **Rust Developers:** `cargo install aichat`\n- **Homebrew/Linuxbrew Users:** `brew install aichat`\n- **Pacman Users**: `pacman -S aichat`\n- **Windows Scoop Users:** `scoop install aichat`\n- **Android Termux Users:** `pkg install aichat`\n\n### Pre-built Binaries\n\nDownload pre-built binaries for macOS, Linux, and Windows from [GitHub Releases](https://github.com/sigoden/aichat/releases), extract them, and add the `aichat` binary to your `$PATH`.\n\n## Features\n\n### Multi-Providers\n\nIntegrate seamlessly with over 20 leading LLM providers through a unified interface. Supported providers include OpenAI, Claude, Gemini (Google AI Studio), Ollama, Groq, Azure-OpenAI, VertexAI, Bedrock, Github Models, Mistral, Deepseek, AI21, XAI Grok, Cohere, Perplexity, Cloudflare, OpenRouter, Ernie, Qianwen, Moonshot, ZhipuAI, MiniMax, Deepinfra, VoyageAI, any OpenAI-Compatible API provider.\n\n### CMD Mode\n\nExplore powerful command-line functionalities with AIChat's CMD mode.\n\n![aichat-cmd](https://github.com/user-attachments/assets/6c58c549-1564-43cf-b772-e1c9fe91d19c)\n\n### REPL Mode\n\nExperience an interactive Chat-REPL with features like tab autocompletion, multi-line input support, history search, configurable keybindings, and custom REPL prompts.\n\n![aichat-repl](https://github.com/user-attachments/assets/218fab08-cdae-4c3b-bcf8-39b6651f1362)\n\n### Shell Assistant\n\nElevate your command-line efficiency. Describe your tasks in natural language, and let AIChat transform them into precise shell commands. AIChat intelligently adjusts to your OS and shell environment.\n\n![aichat-execute](https://github.com/user-attachments/assets/0c77e901-0da2-4151-aefc-a2af96bbb004)\n\n### Multi-Form Input\n\nAccept diverse input forms such as stdin, local files and directories, and remote URLs, allowing flexibility in data handling.\n\n| Input | CMD | REPL |\n| ----------------- | ------------------------------------ | -------------------------------- |\n| CMD | `aichat hello` | |\n| STDIN | `cat data.txt \\| aichat` | |\n| Last Reply | | `.file %%` |\n| Local files | `aichat -f image.png -f data.txt` | `.file image.png data.txt` |\n| Local directories | `aichat -f dir/` | `.file dir/` |\n| Remote URLs | `aichat -f https://example.com` | `.file https://example.com` |\n| External commands | ```aichat -f '`git diff`'``` | ```.file `git diff` ``` |\n| Combine Inputs | `aichat -f dir/ -f data.txt explain` | `.file dir/ data.txt -- explain` |\n\n### Role\n\nCustomize roles to tailor LLM behavior, enhancing interaction efficiency and boosting productivity.\n\n![aichat-role](https://github.com/user-attachments/assets/023df6d2-409c-40bd-ac93-4174fd72f030)\n\n> The role consists of a prompt and model configuration.\n\n### Session\n\nMaintain context-aware conversations through sessions, ensuring continuity in interactions.\n\n![aichat-session](https://github.com/user-attachments/assets/56583566-0f43-435f-95b3-730ae55df031)\n\n> The left side uses a session, while the right side does not use a session.\n\n### Macro\n\nStreamline repetitive tasks by combining a series of REPL commands into a custom macro.\n\n![aichat-macro](https://github.com/user-attachments/assets/23c2a08f-5bd7-4bf3-817c-c484aa74a651)\n\n### RAG\n\nIntegrate external documents into your LLM conversations for more accurate and contextually relevant responses.\n\n![aichat-rag](https://github.com/user-attachments/assets/359f0cb8-ee37-432f-a89f-96a2ebab01f6)\n\n### Function Calling\n\nFunction calling supercharges LLMs by connecting them to external tools and data sources. This unlocks a world of possibilities, enabling LLMs to go beyond their core capabilities and tackle a wider range of tasks.\n\nWe have created a new repository [https://github.com/sigoden/llm-functions](https://github.com/sigoden/llm-functions) to help you make the most of this feature.\n\n#### AI Tools & MCP\n\nIntegrate external tools to automate tasks, retrieve information, and perform actions directly within your workflow.\n\n![aichat-tool](https://github.com/user-attachments/assets/7459a111-7258-4ef0-a2dd-624d0f1b4f92)\n\n#### AI Agents (CLI version of OpenAI GPTs)\n\nAI Agent = Instructions (Prompt) + Tools (Function Callings) + Documents (RAG).\n\n![aichat-agent](https://github.com/user-attachments/assets/0b7e687d-e642-4e8a-b1c1-d2d9b2da2b6b)\n\n### Local Server Capabilities\n\nAIChat includes a lightweight built-in HTTP server for easy deployment.\n\n```\n$ aichat --serve\nChat Completions API: http://127.0.0.1:8000/v1/chat/completions\nEmbeddings API: http://127.0.0.1:8000/v1/embeddings\nRerank API: http://127.0.0.1:8000/v1/rerank\nLLM Playground: http://127.0.0.1:8000/playground\nLLM Arena: http://127.0.0.1:8000/arena?num=2\n```\n\n#### Proxy LLM APIs\n\nThe LLM Arena is a web-based platform where you can compare different LLMs side-by-side. \n\nTest with curl:\n\n```sh\ncurl -X POST -H \"Content-Type: application/json\" -d '{\n \"model\":\"claude:claude-3-5-sonnet-20240620\",\n \"messages\":[{\"role\":\"user\",\"content\":\"hello\"}], \n \"stream\":true\n}' http://127.0.0.1:8000/v1/chat/completions\n```\n\n#### LLM Playground\n\nA web application to interact with supported LLMs directly from your browser.\n\n![aichat-llm-playground](https://github.com/user-attachments/assets/aab1e124-1274-4452-b703-ef15cda55439)\n\n#### LLM Arena\n\nA web platform to compare different LLMs side-by-side.\n\n![aichat-llm-arena](https://github.com/user-attachments/assets/edabba53-a1ef-4817-9153-38542ffbfec6)\n\n## Custom Themes\n\nAIChat supports custom dark and light themes, which highlight response text and code blocks.\n\n![aichat-themes](https://github.com/sigoden/aichat/assets/4012553/29fa8b79-031e-405d-9caa-70d24fa0acf8)\n\n## Documentation\n\n- [Chat-REPL Guide](https://github.com/sigoden/aichat/wiki/Chat-REPL-Guide)\n- [Command-Line Guide](https://github.com/sigoden/aichat/wiki/Command-Line-Guide)\n- [Role Guide](https://github.com/sigoden/aichat/wiki/Role-Guide)\n- [Macro Guide](https://github.com/sigoden/aichat/wiki/Macro-Guide)\n- [RAG Guide](https://github.com/sigoden/aichat/wiki/RAG-Guide)\n- [Environment Variables](https://github.com/sigoden/aichat/wiki/Environment-Variables)\n- [Configuration Guide](https://github.com/sigoden/aichat/wiki/Configuration-Guide)\n- [Custom Theme](https://github.com/sigoden/aichat/wiki/Custom-Theme)\n- [Custom REPL Prompt](https://github.com/sigoden/aichat/wiki/Custom-REPL-Prompt)\n- [FAQ](https://github.com/sigoden/aichat/wiki/FAQ)\n\n## License\n\nCopyright (c) 2023-2025 aichat-developers.\n\nAIChat is made available under the terms of either the MIT License or the Apache License 2.0, at your option.\n\nSee the LICENSE-APACHE and LICENSE-MIT files for license details." }, { "path": "assets/arena.html", "content": "\n\n\n\n \n \n \n AIChat LLM Arena\n \n \n \n \n \n \n \n \n \n \n\n\n\n
\n
\n \n
\n
\n \n \n
handleScrollToBottom(index)\">\n \n \n \n
\n \n \n \n
\n
\n \n
0\">\n \n
\n \n \n
\n
\n
\n \n \n\n\n" }, { "path": "assets/playground.html", "content": "\n\n\n\n \n \n \n AIChat LLM Playground\n \n \n \n \n \n \n \n \n \n \n\n\n\n
\n
\n
\n
\n
AIChat
\n
All-in-one AI-Powered Chat
\n
\n
\n
\n \n \n \n \n
\n
\n \n \n \n \n \n
\n
\n \n \n \n
\n
\n
\n
\n
\n \n \n
\n\n
\n \n \n
\n\n
\n \n \n
\n\n
\n \n \n
\n\n
\n \n \n
\n\n
\n \n \n
\n\n
\n
\n
\n
\n \n
\n
\n \n \n \n
\n
\n
\n
\n \n \n \n \n
\n \n
\n
\n \n \n
\n \n \n \n
\n
\n
\n \n
0\">\n \n
\n \n \n
\n
\n \n
\n
\n
\n
Sessions
\n \n
\n
\n \n
\n
\n
\n
\n \n \n\n\n" }, { "path": "assets/roles/%code%.md", "content": "Provide only code without comments or explanations.\n### INPUT:\nasync sleep in js\n### OUTPUT:\n```javascript\nasync function timeout(ms) {\n return new Promise(resolve => setTimeout(resolve, ms));\n}\n```\n" }, { "path": "assets/roles/%create-prompt%.md", "content": "As a professional Prompt Engineer, your role is to create effective and innovative prompts for interacting with AI models.\n\nYour core skills include:\n1. **CO-STAR Framework Application**: Utilize the CO-STAR framework to build efficient prompts, ensuring effective communication with large language models.\n2. **Contextual Awareness**: Construct prompts that adapt to complex conversation contexts, ensuring relevant and coherent responses.\n3. **Chain-of-Thought Prompting**: Create prompts that elicit AI models to demonstrate their reasoning process, enhancing the transparency and accuracy of answers.\n4. **Zero-shot Learning**: Design prompts that enable AI models to perform specific tasks without requiring examples, reducing dependence on training data.\n5. **Few-shot Learning**: Guide AI models to quickly learn and execute new tasks through a few examples.\n\nYour output format should include:\n- **Context**: Provide comprehensive background information for the task to ensure the AI understands the specific scenario and offers relevant feedback.\n- **Objective**: Clearly define the task objective, guiding the AI to focus on achieving specific goals.\n- **Style**: Specify writing styles according to requirements, such as imitating a particular person or industry expert.\n- **Tone**: Set an appropriate emotional tone to ensure the AI's response aligns with the expected emotional context.\n- **Audience**: Tailor AI responses for a specific audience, ensuring content appropriateness and ease of understanding.\n- **Response**: Specify output formats for easy execution of downstream tasks, such as lists, JSON, or professional reports.\n- **Workflow**: Instruct the AI on how to step-by-step complete tasks, clarifying inputs, outputs, and specific actions for each step.\n- **Examples**: Show a case of input and output that fits the scenario.\n\nYour workflow should be:\n1. Extract key information from user requests to determine design objectives.\n2. Based on user needs, create prompts that meet requirements, with each part being professional and detailed.\n3. Must only output the newly generated and optimized prompts, without explanation, without wrapping it in markdown code block.\n\nMy first request is: __INPUT__\n" }, { "path": "assets/roles/%create-title%.md", "content": "Create a concise, 3-6 word title.\n\n**Notes**:\n- Avoid quotation marks or emojis\n- RESPOND ONLY WITH TITLE SLUG TEXT\n\n**Examples**:\nstock-market-trends\nperfect-chocolate-chip-recipe\nremote-work-productivity-tips\nvideo-game-development-insights\n" }, { "path": "assets/roles/%explain-shell%.md", "content": "Provide a terse, single sentence description of the given shell command.\nDescribe each argument and option of the command.\nProvide short responses in about 80 words.\nAPPLY MARKDOWN formatting when possible." }, { "path": "assets/roles/%functions%.md", "content": "---\nuse_tools: all\n---\n" }, { "path": "assets/roles/%shell%.md", "content": "Provide only {{__shell__}} commands for {{__os_distro__}} without any description.\nEnsure the output is a valid {{__shell__}} command.\nIf there is a lack of details, provide most logical solution.\nIf multiple steps are required, try to combine them using '&&' (For PowerShell, use ';' instead).\nOutput only plain text without any markdown formatting.\n" }, { "path": "config.agent.example.yaml", "content": "# Agent-specific configuration\n# Location `/agents//config.yaml`\n\nmodel: openai:gpt-4o # Specify the LLM to use\ntemperature: null # Set default temperature parameter, range (0, 1)\ntop_p: null # Set default top-p parameter, with a range of (0, 1) or (0, 2) depending on the model\nuse_tools: null # Which additional tools to use by agent. (e.g. 'fs,web_search')\nagent_prelude: null # Set a session to use when starting the agent. (e.g. temp, default)\ninstructions: null # Override the instructions for the agent, have no effect for dynamic instructions\nvariables: # Custom default values for the agent variables\n : \n" }, { "path": "config.example.yaml", "content": "# ---- llm ----\nmodel: openai:gpt-4o # Specify the LLM to use\ntemperature: null # Set default temperature parameter (0, 1)\ntop_p: null # Set default top-p parameter, with a range of (0, 1) or (0, 2) depending on the model\n\n# ---- behavior ----\nstream: true # Controls whether to use the stream-style API.\nsave: true # Indicates whether to persist the message\nkeybindings: emacs # Choose keybinding style (emacs, vi)\neditor: null # Specifies the command used to edit input buffer or session. (e.g. vim, emacs, nano).\nwrap: no # Controls text wrapping (no, auto, )\nwrap_code: false # Enables or disables wrapping of code blocks\n\n# ---- function-calling ----\n# Visit https://github.com/sigoden/llm-functions for setup instructions\nfunction_calling: true # Enables or disables function calling (Globally).\nmapping_tools: # Alias for a tool or toolset\n fs: 'fs_cat,fs_ls,fs_mkdir,fs_rm,fs_write'\nuse_tools: null # Which tools to use by default. (e.g. 'fs,web_search')\n\n# ---- prelude ----\nrepl_prelude: null # Set a default role or session for REPL mode (e.g. role:, session:, :)\ncmd_prelude: null # Set a default role or session for CMD mode (e.g. role:, session:, :)\nagent_prelude: null # Set a session to use when starting a agent (e.g. temp, default)\n\n# ---- session ----\n# Controls the persistence of the session. if true, auto save; if false, not save; if null, asking the user\nsave_session: null\n# Compress session when token count reaches or exceeds this threshold\ncompress_threshold: 4000\n# Text prompt used for creating a concise summary of session message\nsummarize_prompt: 'Summarize the discussion briefly in 200 words or less to use as a prompt for future context.'\n# Text prompt used for including the summary of the entire session\nsummary_prompt: 'This is a summary of the chat history as a recap: '\n\n# ---- RAG ----\n# See [RAG-Guide](https://github.com/sigoden/aichat/wiki/RAG-Guide) for more details.\nrag_embedding_model: null # Specifies the embedding model used for context retrieval\nrag_reranker_model: null # Specifies the reranker model used for sorting retrieved documents\nrag_top_k: 5 # Specifies the number of documents to retrieve for answering queries\nrag_chunk_size: null # Defines the size of chunks for document processing in characters\nrag_chunk_overlap: null # Defines the overlap between chunks\n# Defines the query structure using variables like __CONTEXT__ and __INPUT__ to tailor searches to specific needs\nrag_template: |\n Answer the query based on the context while respecting the rules. (user query, some textual context and rules, all inside xml tags)\n\n \n __CONTEXT__\n \n\n \n - If you don't know, just say so.\n - If you are not sure, ask for clarification.\n - Answer in the same language as the user query.\n - If the context appears unreadable or of poor quality, tell the user then answer as best as you can.\n - If the answer is not in the context but you think you know the answer, explain that to the user then answer with your own knowledge.\n - Answer directly and without using xml tags.\n \n\n \n __INPUT__\n \n\n# Define document loaders to control how RAG and `.file`/`--file` load files of specific formats.\ndocument_loaders:\n # You can add custom loaders using the following syntax:\n # : \n # Note: Use `$1` for input file and `$2` for output file. If `$2` is omitted, use stdout as output.\n pdf: 'pdftotext $1 -' # Load .pdf file, see https://poppler.freedesktop.org to set up pdftotext\n docx: 'pandoc --to plain $1' # Load .docx file, see https://pandoc.org to set up pandoc\n\n# ---- apperence ----\nhighlight: true # Controls syntax highlighting\nlight_theme: false # Activates a light color theme when true. env: AICHAT_LIGHT_THEME\n# Custom REPL left/right prompts, see https://github.com/sigoden/aichat/wiki/Custom-REPL-Prompt for more details\nleft_prompt:\n '{color.green}{?session {?agent {agent}>}{session}{?role /}}{!session {?agent {agent}>}}{role}{?rag @{rag}}{color.cyan}{?session )}{!session >}{color.reset} '\nright_prompt:\n '{color.purple}{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}'\n\n# ---- misc ----\nserve_addr: 127.0.0.1:8000 # Server listening address \nuser_agent: null # Set User-Agent HTTP header, use `auto` for aichat/\nsave_shell_history: true # Whether to save shell execution command to the history file\n# URL to sync model changes from, e.g., https://cdn.jsdelivr.net/gh/sigoden/aichat@main/models.yaml\nsync_models_url: https://raw.githubusercontent.com/sigoden/aichat/refs/heads/main/models.yaml\n\n# ---- clients ----\nclients:\n # All clients have the following configuration:\n # - type: xxxx\n # name: xxxx # Only use it to distinguish clients with the same client type. Optional\n # models:\n # - name: xxxx # Chat model\n # max_input_tokens: 100000\n # supports_vision: true\n # supports_function_calling: true\n # - name: xxxx # Embedding model\n # type: embedding\n # default_chunk_size: 1500 \n # max_batch_size: 100\n # - name: xxxx # Reranker model\n # type: reranker \n # patch: # Patch api\n # chat_completions: # Api type, possible values: chat_completions, embeddings, and rerank\n # : # The regex to match model names, e.g. '.*' 'gpt-4o' 'gpt-4o|gpt-4-.*'\n # url: '' # Patch request url\n # body: # Patch request body\n # \n # headers: # Patch request headers\n # : \n # extra:\n # proxy: socks5://127.0.0.1:1080 # Set proxy\n # connect_timeout: 10 # Set timeout in seconds for connect to api\n\n # See https://platform.openai.com/docs/quickstart\n - type: openai\n api_base: https://api.openai.com/v1 # Optional\n api_key: xxx\n organization_id: org-xxx # Optional\n\n # For any platform compatible with OpenAI's API\n - type: openai-compatible\n name: ollama\n api_base: http://localhost:11434/v1\n api_key: xxx # Optional\n models:\n - name: deepseek-r1\n max_input_tokens: 131072\n - name: llama3.1\n max_input_tokens: 128000\n supports_function_calling: true\n - name: llama3.2-vision\n max_input_tokens: 131072\n supports_vision: true\n - name: nomic-embed-text\n type: embedding\n default_chunk_size: 1000\n max_batch_size: 50\n\n # See https://ai.google.dev/docs\n - type: gemini\n api_base: https://generativelanguage.googleapis.com/v1beta\n api_key: xxx\n patch:\n chat_completions:\n '.*':\n body:\n safetySettings:\n - category: HARM_CATEGORY_HARASSMENT\n threshold: BLOCK_NONE\n - category: HARM_CATEGORY_HATE_SPEECH\n threshold: BLOCK_NONE\n - category: HARM_CATEGORY_SEXUALLY_EXPLICIT\n threshold: BLOCK_NONE\n - category: HARM_CATEGORY_DANGEROUS_CONTENT\n threshold: BLOCK_NONE\n\n # See https://docs.anthropic.com/claude/reference/getting-started-with-the-api\n - type: claude\n api_base: https://api.anthropic.com/v1 # Optional\n api_key: xxx\n\n # See https://docs.mistral.ai/\n - type: openai-compatible\n name: mistral\n api_base: https://api.mistral.ai/v1\n api_key: xxx\n\n # See https://docs.x.ai/docs\n - type: openai-compatible\n name: xai\n api_base: https://api.x.ai/v1\n api_key: xxx\n\n # See https://docs.ai21.com/docs/quickstart\n - type: openai-compatible\n name: ai12\n api_base: https://api.ai21.com/studio/v1\n api_key: xxx\n\n # See https://docs.cohere.com/docs/the-cohere-platform\n - type: cohere\n api_base: https://api.cohere.ai/v2 # Optional\n api_key: xxx\n\n # See https://docs.perplexity.ai/docs/getting-started\n - type: openai-compatible\n name: perplexity\n api_base: https://api.perplexity.ai\n api_key: xxx\n\n # See https://console.groq.com/docs/quickstart\n - type: openai-compatible\n name: groq\n api_base: https://api.groq.com/openai/v1\n api_key: xxx\n\n # See https://learn.microsoft.com/en-us/azure/ai-services/openai/chatgpt-quickstart\n - type: azure-openai\n api_base: https://{RESOURCE}.openai.azure.com\n api_key: xxx\n models:\n - name: gpt-4o # Model deployment name\n max_input_tokens: 128000\n supports_vision: true\n supports_function_calling: true\n\n # See https://cloud.google.com/vertex-ai\n - type: vertexai\n project_id: xxx\n location: xxx\n # Specifies a application-default-credentials (adc) file\n # Run `gcloud auth application-default login` to init the adc file\n # see https://cloud.google.com/docs/authentication/external/set-up-adc\n adc_file: /application_default_credentials.json> # Optional field\n patch:\n chat_completions:\n 'gemini-.*':\n body:\n safetySettings:\n - category: HARM_CATEGORY_HARASSMENT\n threshold: BLOCK_ONLY_HIGH\n - category: HARM_CATEGORY_HATE_SPEECH\n threshold: BLOCK_ONLY_HIGH\n - category: HARM_CATEGORY_SEXUALLY_EXPLICIT\n threshold: BLOCK_ONLY_HIGH\n - category: HARM_CATEGORY_DANGEROUS_CONTENT\n threshold: BLOCK_ONLY_HIGH\n\n # See https://docs.aws.amazon.com/bedrock/latest/userguide/\n - type: bedrock\n access_key_id: xxx\n secret_access_key: xxx\n region: xxx\n session_token: xxx # Optional, only needed for temporary credentials\n\n # See https://developers.cloudflare.com/workers-ai/\n - type: openai-compatible\n name: cloudflare\n api_base: https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/v1\n api_key: xxx\n\n # See https://cloud.baidu.com/doc/WENXINWORKSHOP/index.html\n - type: openai-compatible\n name: ernie\n api_base: https://qianfan.baidubce.com/v2\n api_key: xxx\n\n # See https://dashscope.aliyun.com/\n - type: openai-compatible\n name: qianwen\n api_base: https://dashscope.aliyuncs.com/compatible-mode/v1\n api_key: xxx\n\n # See https://cloud.tencent.com/product/hunyuan\n - type: openai-compatible\n name: hunyuan\n api_base: https://api.hunyuan.cloud.tencent.com/v1\n api_key: xxx\n\n # See https://platform.moonshot.cn/docs/intro\n - type: openai-compatible\n name: moonshot\n api_base: https://api.moonshot.cn/v1\n api_key: xxx\n\n # See https://platform.deepseek.com/api-docs/\n - type: openai-compatible\n name: deepseek\n api_base: https://api.deepseek.com\n api_key: xxx\n\n # See https://open.bigmodel.cn/dev/howuse/introduction\n - type: openai-compatible\n name: zhipuai\n api_base: https://open.bigmodel.cn/api/paas/v4\n api_key: xxx\n\n # See https://platform.minimaxi.com/document/Fast%20access\n - type: openai-compatible\n name: minimax\n api_base: https://api.minimax.chat/v1\n api_key: xxx\n\n # See https://openrouter.ai/docs#quick-start\n - type: openai-compatible\n name: openrouter\n api_base: https://openrouter.ai/api/v1\n api_key: xxx\n\n # See https://github.com/marketplace/models\n - type: openai-compatible\n name: github\n api_base: https://models.inference.ai.azure.com\n api_key: xxx\n\n # See https://deepinfra.com/docs\n - type: openai-compatible\n name: deepinfra\n api_base: https://api.deepinfra.com/v1/openai\n api_key: xxx\n\n\n # ----- RAG dedicated -----\n\n # See https://jina.ai\n - type: openai-compatible\n name: jina\n api_base: https://api.jina.ai/v1\n api_key: xxx\n\n # See https://docs.voyageai.com/docs/introduction\n - type: openai-compatible\n name: voyageai\n api_base: https://api.voyageai.com/v1\n api_key: xxx\n" }, { "path": "models.yaml", "content": "# Links:\n# - https://platform.openai.com/docs/models\n# - https://platform.openai.com/docs/api-reference/chat\n- provider: openai\n models:\n - name: gpt-5.2\n max_input_tokens: 400000\n max_output_tokens: 128000\n input_price: 1.75\n output_price: 14\n supports_vision: true\n supports_function_calling: true\n - name: gpt-5\n max_input_tokens: 400000\n max_output_tokens: 128000\n input_price: 1.25\n output_price: 10\n supports_vision: true\n supports_function_calling: true\n - name: gpt-5-mini\n max_input_tokens: 400000\n max_output_tokens: 128000\n input_price: 0.25\n output_price: 2\n supports_vision: true\n supports_function_calling: true\n - name: gpt-5-nano\n max_input_tokens: 400000\n max_output_tokens: 128000\n input_price: 0.05\n output_price: 0.4\n supports_vision: true\n supports_function_calling: true\n - name: gpt-4.1\n max_input_tokens: 1047576\n max_output_tokens: 32768\n input_price: 2\n output_price: 8\n supports_vision: true\n supports_function_calling: true\n - name: gpt-4o\n max_input_tokens: 128000\n max_output_tokens: 16384\n input_price: 2.5\n output_price: 10\n supports_vision: true\n supports_function_calling: true\n - name: gpt-4-turbo\n max_input_tokens: 128000\n max_output_tokens: 4096\n input_price: 10\n output_price: 30\n supports_vision: true\n supports_function_calling: true\n - name: gpt-3.5-turbo\n max_input_tokens: 16385\n max_output_tokens: 4096\n input_price: 0.5\n output_price: 1.5\n supports_function_calling: true\n - name: text-embedding-3-large\n type: embedding\n input_price: 0.13\n max_tokens_per_chunk: 8191\n default_chunk_size: 2000\n max_batch_size: 100\n - name: text-embedding-3-small\n type: embedding\n input_price: 0.02\n max_tokens_per_chunk: 8191\n default_chunk_size: 2000\n max_batch_size: 100\n\n# Links:\n# - https://ai.google.dev/models/gemini\n# - https://ai.google.dev/pricing\n# - https://ai.google.dev/api/rest/v1beta/models/streamGenerateContent\n- provider: gemini\n models:\n - name: gemini-2.5-flash\n max_input_tokens: 1048576\n max_output_tokens: 65536\n input_price: 0\n output_price: 0\n supports_vision: true\n supports_function_calling: true\n - name: gemini-2.5-pro\n max_input_tokens: 1048576\n max_output_tokens: 65536\n input_price: 0\n output_price: 0\n supports_vision: true\n supports_function_calling: true\n - name: gemini-2.5-flash-lite\n max_input_tokens: 1000000\n max_output_tokens: 64000\n input_price: 0\n output_price: 0\n supports_vision: true\n supports_function_calling: true\n - name: gemini-3-pro-preview\n max_input_tokens: 1048576\n supports_vision: true\n supports_function_calling: true\n - name: gemini-3-flash-preview\n max_input_tokens: 1048576\n supports_vision: true\n supports_function_calling: true\n - name: gemini-2.0-flash\n max_input_tokens: 1048576\n max_output_tokens: 8192\n input_price: 0\n output_price: 0\n supports_vision: true\n supports_function_calling: true\n - name: gemini-2.0-flash-lite\n max_input_tokens: 1048576\n max_output_tokens: 8192\n input_price: 0\n output_price: 0\n supports_vision: true\n supports_function_calling: true\n - name: gemma-3-27b-it\n max_input_tokens: 131072\n max_output_tokens: 8192\n input_price: 0\n output_price: 0\n - name: text-embedding-004\n type: embedding\n input_price: 0\n max_tokens_per_chunk: 2048\n default_chunk_size: 1500\n max_batch_size: 100\n\n# Links:\n# - https://docs.anthropic.com/en/docs/about-claude/models/all-models\n# - https://docs.anthropic.com/en/api/messages\n- provider: claude\n models:\n - name: claude-opus-4-6\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n supports_function_calling: true\n - name: claude-opus-4-6:thinking\n real_name: claude-opus-4-6\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n supports_function_calling: true\n patch:\n body:\n temperature: null\n top_p: null\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: claude-sonnet-4-6\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n supports_function_calling: true\n - name: claude-sonnet-4-6:thinking\n real_name: claude-sonnet-4-6\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n supports_function_calling: true\n patch:\n body:\n temperature: null\n top_p: null\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: claude-opus-4-5-20251101\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n supports_function_calling: true\n - name: claude-opus-4-5-20251101:thinking\n real_name: claude-opus-4-5-20251101\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n supports_function_calling: true\n patch:\n body:\n temperature: null\n top_p: null\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: claude-sonnet-4-5-20250929\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n supports_function_calling: true\n - name: claude-sonnet-4-5-20250929:thinking\n real_name: claude-sonnet-4-5-20250929\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n supports_function_calling: true\n patch:\n body:\n temperature: null\n top_p: null\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: claude-haiku-4-5-20251001\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 1\n output_price: 5\n supports_vision: true\n supports_function_calling: true\n - name: claude-haiku-4-5-20251001:thinking\n real_name: claude-haiku-4-5-20251001\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 1\n output_price: 5\n supports_vision: true\n supports_function_calling: true\n patch:\n body:\n temperature: null\n top_p: null\n thinking:\n type: enabled\n budget_tokens: 16000\n\n# Links:\n# - https://docs.mistral.ai/getting-started/models/models_overview/\n# - https://mistral.ai/pricing#api-pricing\n# - https://docs.mistral.ai/api/\n- provider: mistral\n models:\n - name: mistral-large-latest\n max_output_tokens: 262144\n input_price: 0.5\n output_price: 1.5\n supports_function_calling: true\n supports_vision: true\n - name: mistral-medium-latest\n max_input_tokens: 131072\n input_price: 0.4\n output_price: 2\n supports_function_calling: true\n supports_vision: true\n - name: mistral-small-latest\n max_input_tokens: 32768\n input_price: 0.1\n output_price: 0.3\n supports_function_calling: true\n supports_vision: true\n - name: magistral-medium-latest\n max_input_tokens: 131072\n input_price: 2\n output_price: 5\n - name: magistral-small-latest\n max_input_tokens: 131072\n input_price: 0.5\n output_price: 1.5\n - name: devstral-medium-latest\n max_input_tokens: 262144\n input_price: 0.4\n output_price: 2\n supports_function_calling: true\n - name: devstral-small-latest\n max_input_tokens: 262144\n input_price: 0.1\n output_price: 0.3\n supports_function_calling: true\n - name: codestral-latest\n max_input_tokens: 262144\n input_price: 0.3\n output_price: 0.9\n supports_function_calling: true\n - name: ministral-14b-latest\n max_input_tokens: 262144\n input_price: 0.2\n output_price: 0.2\n supports_function_calling: true\n - name: mistral-embed\n type: embedding\n max_input_tokens: 8092\n input_price: 0.1\n max_tokens_per_chunk: 8092\n default_chunk_size: 2000\n\n# Links:\n# - https://docs.ai21.com/docs/jamba-foundation-models\n# - https://www.ai21.com/pricing\n# - https://docs.ai21.com/reference/jamba-1-6-api-ref\n- provider: ai21\n models:\n - name: jamba-large\n max_input_tokens: 256000\n input_price: 2\n output_price: 8\n supports_function_calling: true\n - name: jamba-mini\n max_input_tokens: 256000\n input_price: 0.2\n output_price: 0.4\n supports_function_calling: true\n\n# Links:\n# - https://docs.cohere.com/docs/models\n# - https://cohere.com/pricing\n# - https://docs.cohere.com/reference/chat\n- provider: cohere\n models:\n - name: command-a-03-2025\n max_input_tokens: 262144\n max_output_tokens: 8192\n input_price: 2.5\n output_price: 10\n supports_function_calling: true\n - name: command-a-reasoning-08-2025\n max_input_tokens: 262144\n max_output_tokens: 32768\n input_price: 2.5\n output_price: 10\n - name: command-a-vision-07-2025\n max_input_tokens: 131072\n max_output_tokens: 8192\n input_price: 2.5\n output_price: 10\n supports_vision: true\n - name: command-r7b-12-2024\n max_input_tokens: 131072\n max_output_tokens: 4096\n input_price: 0.0375\n output_price: 0.15\n - name: embed-v4.0\n type: embedding\n input_price: 0.12\n max_tokens_per_chunk: 2048\n default_chunk_size: 2000\n max_batch_size: 96\n - name: embed-english-v3.0\n type: embedding\n input_price: 0.1\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 96\n - name: embed-multilingual-v3.0\n type: embedding\n input_price: 0.1\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 96\n - name: rerank-v3.5\n type: reranker\n max_input_tokens: 4096\n - name: rerank-english-v3.0\n type: reranker\n max_input_tokens: 4096\n - name: rerank-multilingual-v3.0\n type: reranker\n max_input_tokens: 4096\n\n# Links:\n# - https://docs.x.ai/docs/models\n# - https://docs.x.ai/docs/api-reference#chat-completions\n- provider: xai\n models:\n - name: grok-4-1-fast-non-reasoning\n max_input_tokens: 2000000\n input_price: 0.2\n output_price: 0.5\n supports_function_calling: true\n - name: grok-4-1-fast-reasoning\n max_input_tokens: 2000000\n input_price: 0.2\n output_price: 0.5\n supports_function_calling: true\n - name: grok-code-fast-1\n max_input_tokens: 256000\n input_price: 0.2\n output_price: 1.5\n supports_function_calling: true\n\n# Links:\n# - https://docs.perplexity.ai/getting-started/models\n# - https://docs.perplexity.ai/api-reference/chat-completions\n- provider: perplexity\n models:\n - name: sonar-pro\n max_input_tokens: 200000\n input_price: 3\n output_price: 15\n - name: sonar\n max_input_tokens: 128000\n input_price: 1\n output_price: 1\n - name: sonar-reasoning-pro\n max_input_tokens: 128000\n input_price: 2\n output_price: 8\n - name: sonar-deep-research\n max_input_tokens: 128000\n input_price: 2\n output_price: 8\n\n# Links:\n# - https://console.groq.com/docs/models\n# - https://console.groq.com/docs/api-reference#chat\n- provider: groq\n models:\n - name: openai/gpt-oss-120b\n max_input_tokens: 131072\n input_price: 0\n output_price: 0\n supports_function_calling: true\n - name: openai/gpt-oss-20b\n max_input_tokens: 131072\n input_price: 0\n output_price: 0\n supports_function_calling: true\n - name: meta-llama/llama-4-maverick-17b-128e-instruct\n max_input_tokens: 131072\n input_price: 0\n output_price: 0\n supports_vision: true\n supports_function_calling: true\n - name: meta-llama/llama-4-scout-17b-16e-instruct\n max_input_tokens: 131072\n input_price: 0\n output_price: 0\n supports_vision: true\n supports_function_calling: true\n - name: llama-3.3-70b-versatile\n max_input_tokens: 131072\n input_price: 0\n output_price: 0\n supports_function_calling: true\n - name: moonshotai/kimi-k2-instruct-0905\n max_input_tokens: 262144\n input_price: 0\n output_price: 0\n supports_function_calling: true\n - name: qwen/qwen3-32b\n max_input_tokens: 131072\n input_price: 0\n output_price: 0\n - name: groq/compound\n max_input_tokens: 131072\n input_price: 0\n output_price: 0\n - name: groq/compound-mini\n max_input_tokens: 131072\n input_price: 0\n output_price: 0\n\n# Links:\n# - https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models\n# - https://cloud.google.com/vertex-ai/generative-ai/pricing\n# - https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/gemini\n- provider: vertexai\n models:\n - name: gemini-2.5-flash\n max_input_tokens: 1048576\n max_output_tokens: 65536\n input_price: 0.3\n output_price: 2.5\n supports_vision: true\n supports_function_calling: true\n - name: gemini-2.5-pro\n max_input_tokens: 1048576\n max_output_tokens: 65536\n input_price: 1.25\n output_price: 10\n supports_vision: true\n supports_function_calling: true\n - name: gemini-2.5-flash-lite\n max_input_tokens: 1048576\n max_output_tokens: 65536\n input_price: 0.3\n output_price: 0.4\n supports_vision: true\n supports_function_calling: true\n - name: gemini-3-pro-preview\n max_input_tokens: 1048576\n supports_vision: true\n supports_function_calling: true\n - name: gemini-3-flash-preview\n max_input_tokens: 1048576\n supports_vision: true\n supports_function_calling: true\n - name: gemini-2.0-flash-001\n max_input_tokens: 1048576\n max_output_tokens: 8192\n input_price: 0.15\n output_price: 0.6\n supports_vision: true\n supports_function_calling: true\n - name: gemini-2.0-flash-lite-001\n max_input_tokens: 1048576\n max_output_tokens: 8192\n input_price: 0.075\n output_price: 0.3\n supports_vision: true\n supports_function_calling: true\n - name: claude-opus-4-6\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n supports_function_calling: true\n - name: claude-opus-4-6:thinking\n real_name: claude-opus-4-6\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n patch:\n body:\n temperature: null\n top_p: null\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: claude-sonnet-4-6\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n supports_function_calling: true\n - name: claude-sonnet-4-6:thinking\n real_name: claude-sonnet-4-6\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n patch:\n body:\n temperature: null\n top_p: null\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: claude-opus-4-5@20251101\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n supports_function_calling: true\n - name: claude-opus-4-5@20251101:thinking\n real_name: claude-opus-4-5@20251101\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n patch:\n body:\n temperature: null\n top_p: null\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: claude-sonnet-4-5@20250929\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n supports_function_calling: true\n - name: claude-sonnet-4-5@20250929:thinking\n real_name: claude-sonnet-4-5@20250929\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n patch:\n body:\n temperature: null\n top_p: null\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: claude-haiku-4-5@20251001\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 1\n output_price: 5\n supports_vision: true\n supports_function_calling: true\n - name: claude-haiku-4-5@20251001:thinking\n real_name: claude-haiku-4-5@20251001\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 1\n output_price: 5\n supports_vision: true\n patch:\n body:\n temperature: null\n top_p: null\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: text-embedding-005\n type: embedding\n max_input_tokens: 20000\n input_price: 0.025\n max_tokens_per_chunk: 2048\n default_chunk_size: 1500\n max_batch_size: 5\n - name: text-multilingual-embedding-002\n type: embedding\n max_input_tokens: 20000\n input_price: 0.2\n max_tokens_per_chunk: 2048\n default_chunk_size: 1500\n max_batch_size: 5\n\n# Links:\n# - https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html#model-ids-arns\n# - https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-supported-models-features.html\n# - https://aws.amazon.com/bedrock/pricing/\n# - https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-call.html\n- provider: bedrock\n models:\n - name: us.anthropic.claude-opus-4-6-v1\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n supports_function_calling: true\n - name: us.anthropic.claude-opus-4-6-v1:thinking\n real_name: us.anthropic.claude-opus-4-6-v1\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n patch:\n body:\n inferenceConfig:\n temperature: null\n topP: null\n additionalModelRequestFields:\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: us.anthropic.claude-sonnet-4-6\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n supports_function_calling: true\n - name: us.anthropic.claude-sonnet-4-6:thinking\n real_name: us.anthropic.claude-sonnet-4-6\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n patch:\n body:\n inferenceConfig:\n temperature: null\n topP: null\n additionalModelRequestFields:\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: us.anthropic.claude-opus-4-5-20251101-v1:0\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n supports_function_calling: true\n - name: us.anthropic.claude-opus-4-5-20251101-v1:0:thinking\n real_name: us.anthropic.claude-opus-4-5-20251101-v1:0\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n patch:\n body:\n inferenceConfig:\n temperature: null\n topP: null\n additionalModelRequestFields:\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: us.anthropic.claude-sonnet-4-5-20250929-v1:0\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n supports_function_calling: true\n - name: us.anthropic.claude-sonnet-4-5-20250929-v1:0:thinking\n real_name: us.anthropic.claude-sonnet-4-5-20250929-v1:0\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n patch:\n body:\n inferenceConfig:\n temperature: null\n topP: null\n additionalModelRequestFields:\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: us.anthropic.claude-haiku-4-5-20251001-v1:0\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 1\n output_price: 5\n supports_vision: true\n supports_function_calling: true\n - name: us.anthropic.claude-haiku-4-5-20251001-v1:0:thinking\n real_name: us.anthropic.claude-haiku-4-5-20251001-v1:0\n max_input_tokens: 200000\n max_output_tokens: 24000\n require_max_tokens: true\n input_price: 1\n output_price: 5\n supports_vision: true\n patch:\n body:\n inferenceConfig:\n temperature: null\n topP: null\n additionalModelRequestFields:\n thinking:\n type: enabled\n budget_tokens: 16000\n - name: us.meta.llama4-maverick-17b-instruct-v1:0\n max_input_tokens: 131072\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 0.24\n output_price: 0.97\n supports_function_calling: true\n supports_vision: true\n - name: us.meta.llama4-scout-17b-instruct-v1:0\n max_input_tokens: 131072\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 0.17\n output_price: 0.66\n supports_function_calling: true\n supports_vision: true\n - name: us.meta.llama3-3-70b-instruct-v1:0\n max_input_tokens: 131072\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 0.72\n output_price: 0.72\n supports_function_calling: true\n - name: us.amazon.nova-premier-v1:0\n max_input_tokens: 300000\n max_output_tokens: 5120\n input_price: 2.5\n output_price: 12.5\n - name: us.amazon.nova-pro-v1:0\n max_input_tokens: 300000\n max_output_tokens: 5120\n input_price: 0.8\n output_price: 3.2\n supports_vision: true\n - name: us.amazon.nova-lite-v1:0\n max_input_tokens: 300000\n max_output_tokens: 5120\n input_price: 0.06\n output_price: 0.24\n supports_vision: true\n - name: us.amazon.nova-micro-v1:0\n max_input_tokens: 128000\n max_output_tokens: 5120\n input_price: 0.035\n output_price: 0.14\n - name: cohere.embed-english-v3\n type: embedding\n input_price: 0.1\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 96\n - name: cohere.embed-multilingual-v3\n type: embedding\n input_price: 0.1\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 96\n - name: us.deepseek.r1-v1:0\n max_input_tokens: 128000\n input_price: 1.35\n output_price: 5.4\n\n# Links:\n# - https://developers.cloudflare.com/workers-ai/models/\n# - https://developers.cloudflare.com/workers-ai/configuration/open-ai-compatibility/\n- provider: cloudflare\n models:\n - name: '@cf/meta/llama-4-scout-17b-16e-instruct'\n max_input_tokens: 131072\n max_output_tokens: 2048\n require_max_tokens: true\n input_price: 0\n output_price: 0\n - name: '@cf/meta/llama-3.3-70b-instruct-fp8-fast'\n max_input_tokens: 131072\n max_output_tokens: 2048\n require_max_tokens: true\n input_price: 0\n output_price: 0\n - name: '@cf/qwen/qwen3-30b-a3b-fp8'\n max_input_tokens: 131072\n max_output_tokens: 2048\n require_max_tokens: true\n input_price: 0\n output_price: 0\n - name: '@cf/qwen/qwen2.5-coder-32b-instruct'\n max_input_tokens: 131072\n max_output_tokens: 2048\n require_max_tokens: true\n input_price: 0\n output_price: 0\n - name: '@cf/zai-org/glm-4.7-flash'\n max_input_tokens: 131072\n max_output_tokens: 2048\n require_max_tokens: true\n input_price: 0\n output_price: 0\n - name: '@cf/google/gemma-3-12b-it'\n max_input_tokens: 131072\n max_output_tokens: 2048\n require_max_tokens: true\n input_price: 0\n output_price: 0\n - name: '@cf/mistralai/mistral-small-3.1-24b-instruct'\n max_input_tokens: 131072\n max_output_tokens: 2048\n require_max_tokens: true\n input_price: 0\n output_price: 0\n - name: '@cf/baai/bge-large-en-v1.5'\n type: embedding\n input_price: 0\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 100\n\n# Links:\n# - https://cloud.baidu.com/doc/qianfan/s/rmh4stp0j\n# - https://cloud.baidu.com/doc/qianfan/s/wmh4sv6ya\n- provider: ernie\n models:\n - name: ernie-4.5-turbo-128k\n max_input_tokens: 131072\n input_price: 0.112\n output_price: 0.448\n - name: ernie-4.5-turbo-vl-32k\n max_input_tokens: 32768\n input_price: 0.42\n output_price: 1.26\n supports_vision: true\n - name: ernie-5.0-thinking-preview\n max_input_tokens: 131072\n input_price: 1.4\n output_price: 5.6\n - name: ernie-x1.1-preview\n max_input_tokens: 65536\n input_price: 0.14\n output_price: 0.56\n - name: bge-large-zh\n type: embedding\n input_price: 0.07\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 16\n - name: bge-large-en\n type: embedding\n input_price: 0.07\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 16\n - name: bce-reranker-base\n type: reranker\n max_input_tokens: 1024\n input_price: 0.07\n\n\n# Links:\n# - https://help.aliyun.com/zh/model-studio/getting-started/models\n# - https://help.aliyun.com/zh/model-studio/developer-reference/use-qwen-by-calling-api\n- provider: qianwen\n models:\n - name: qwen3.5-plus\n max_input_tokens: 262144\n supports_function_calling: true\n patch:\n body:\n enable_thinking: false\n - name: qwen3.5-plus:thinking\n real_name: qwen3.5-plus\n max_input_tokens: 262144\n supports_function_calling: true\n - name: qwen3-max\n max_input_tokens: 262144\n supports_function_calling: true\n - name: qwen3-max:thinking\n real_name: qwen3-max\n max_input_tokens: 262144\n supports_function_calling: true\n patch:\n body:\n enable_thinking: true\n - name: qwen3-vl-plus\n max_input_tokens: 262144\n supports_vision: true\n - name: qwen3-vl-flash\n max_input_tokens: 262144\n supports_vision: true\n - name: qwen3-coder-plus\n max_input_tokens: 1000000\n - name: qwen3-coder-flash\n max_input_tokens: 1000000\n - name: qwen3.5-397b-a17b\n max_input_tokens: 262144\n supports_function_calling: true\n patch:\n body:\n enable_thinking: false\n - name: qwen3.5-397b-a17b:thinking\n real_name: qwen3.5-397b-a17b\n max_input_tokens: 262144\n supports_function_calling: true\n - name: qwen3-next-80b-a3b-instruct\n max_input_tokens: 131072\n input_price: 0.14\n output_price: 0.56\n supports_function_calling: true\n - name: qwen3-next-80b-a3b-thinking\n max_input_tokens: 131072\n input_price: 0.14\n output_price: 1.4\n - name: qwen3-235b-a22b-instruct-2507\n max_input_tokens: 131072\n input_price: 0.28\n output_price: 1.12\n supports_function_calling: true\n - name: qwen3-235b-a22b-thinking-2507\n max_input_tokens: 131072\n input_price: 0.28\n output_price: 2.8\n - name: qwen3-30b-a3b-instruct-2507\n max_input_tokens: 131072\n input_price: 0.105\n output_price: 0.42\n supports_function_calling: true\n - name: qwen3-30b-a3b-thinking-2507\n max_input_tokens: 131072\n input_price: 0.105\n output_price: 1.05 \n - name: qwen3-vl-32b-instruct\n max_input_tokens: 131072\n input_price: 0.28\n output_price: 1.12\n supports_vision: true\n - name: qwen3-vl-8b-instruct\n max_input_tokens: 131072\n input_price: 0.07\n output_price: 0.28\n supports_vision: true\n - name: qwen3-coder-next\n max_input_tokens: 262144\n - name: qwen3-coder-480b-a35b-instruct\n max_input_tokens: 262144\n - name: qwen3-coder-30b-a3b-instruct\n max_input_tokens: 262144\n - name: text-embedding-v4\n type: embedding\n input_price: 0.1\n max_tokens_per_chunk: 8192\n default_chunk_size: 2000\n max_batch_size: 10\n - name: text-embedding-v3\n type: embedding\n input_price: 0.1\n max_tokens_per_chunk: 8192\n default_chunk_size: 2000\n max_batch_size: 10\n\n# links:\n# - https://cloud.tencent.com/document/product/1729/104753\n# - https://cloud.tencent.com/document/product/1729/97731\n# - https://cloud.tencent.com/document/product/1729/111007\n- provider: hunyuan\n models:\n - name: hunyuan-2.0-instruct-20251111\n max_input_tokens: 131072\n input_price: 0.112\n output_price: 0.28\n supports_function_calling: true\n - name: hunyuan-2.0-thinking-20251109\n max_input_tokens: 131072\n input_price: 0.14\n output_price: 0.56\n supports_function_calling: true\n - name: hunyuan-vision-1.5-instruct\n max_input_tokens: 24576\n input_price: 0.42\n output_price: 1.26\n supports_vision: true\n - name: hunyuan-embedding\n type: embedding\n input_price: 0.01\n max_tokens_per_chunk: 1024\n default_chunk_size: 1000\n max_batch_size: 100\n\n# Links:\n# - https://platform.moonshot.cn/docs/pricing/chat#%E8%AE%A1%E8%B4%B9%E5%9F%BA%E6%9C%AC%E6%A6%82%E5%BF%B5\n# - https://platform.moonshot.cn/docs/api/chat#%E5%85%AC%E5%BC%80%E7%9A%84%E6%9C%8D%E5%8A%A1%E5%9C%B0%E5%9D%80\n- provider: moonshot\n models:\n - name: kimi-k2.5\n max_input_tokens: 262144\n input_price: 0.56\n output_price: 2.94\n supports_vision: true\n supports_function_calling: true\n - name: kimi-k2-turbo-preview\n max_input_tokens: 262144\n input_price: 1.12\n output_price: 8.12\n supports_vision: true\n supports_function_calling: true\n - name: kimi-k2-0905-preview\n max_input_tokens: 262144\n input_price: 0.56\n output_price: 2.24\n supports_vision: true\n supports_function_calling: true\n - name: kimi-k2-thinking-turbo\n max_input_tokens: 262144\n input_price: 1.12\n output_price: 8.12\n supports_vision: true\n - name: kimi-k2-thinking\n max_input_tokens: 262144\n input_price: 0.56\n output_price: 2.24\n supports_vision: true\n\n# Links:\n# - https://api-docs.deepseek.com/quick_start/pricing\n# - https://platform.deepseek.com/api-docs/api/create-chat-completion\n- provider: deepseek\n models:\n - name: deepseek-chat\n max_input_tokens: 64000\n max_output_tokens: 8192\n input_price: 0.56\n output_price: 1.68\n supports_function_calling: true\n - name: deepseek-reasoner\n max_input_tokens: 64000\n max_output_tokens: 32768\n input_price: 0.56\n output_price: 1.68\n\n# Links:\n# - https://open.bigmodel.cn/pricing\n# - https://open.bigmodel.cn/dev/api#glm-4\n- provider: zhipuai\n models:\n - name: glm-5\n max_input_tokens: 202752\n supports_function_calling: true\n - name: glm-5:instruct\n real_name: glm-5\n max_input_tokens: 202752\n supports_function_calling: true\n patch:\n body:\n thinking:\n type: disabled\n - name: glm-4.7\n max_input_tokens: 202752\n supports_function_calling: true\n - name: glm-4.7:instruct\n real_name: glm-4.7\n max_input_tokens: 202752\n supports_function_calling: true\n patch:\n body:\n thinking:\n type: disabled\n - name: glm-4.7-flash\n max_input_tokens: 202752\n input_price: 0\n output_price: 0\n supports_function_calling: true\n - name: glm-4.6v\n max_input_tokens: 65536\n supports_vision: true\n - name: glm-4.6v-flash\n max_input_tokens: 65536\n input_price: 0\n output_price: 0\n supports_vision: true\n - name: embedding-3\n type: embedding\n max_input_tokens: 8192\n input_price: 0.07\n max_tokens_per_chunk: 8192\n default_chunk_size: 2000\n - name: rerank\n type: reranker\n max_input_tokens: 4096\n input_price: 0.112\n\n# Links:\n# - https://platform.minimaxi.com/docs/guides/pricing-paygo\n# - https://platform.minimaxi.com/document/ChatCompletion%20v2\n- provider: minimax\n models:\n - name: minimax-m2.5\n max_input_tokens: 204800\n input_price: 0.294\n output_price: 1.176\n supports_function_calling: true\n - name: minimax-m2.5-highspeed\n max_input_tokens: 204800\n input_price: 0.588\n output_price: 2.352\n supports_function_calling: true\n - name: minimax-m2.1\n max_input_tokens: 204800\n input_price: 0.294\n output_price: 1.176\n supports_function_calling: true\n - name: minimax-m2.1-highspeed\n max_input_tokens: 204800\n input_price: 0.588\n output_price: 2.352\n supports_function_calling: true\n\n# Links:\n# - https://openrouter.ai/models\n# - https://openrouter.ai/docs/api-reference/chat-completion\n- provider: openrouter\n models:\n - name: openai/gpt-5.2\n max_input_tokens: 400000\n max_output_tokens: 128000\n input_price: 1.75\n output_price: 14\n supports_vision: true\n supports_function_calling: true\n - name: openai/gpt-5\n max_input_tokens: 400000\n max_output_tokens: 128000\n input_price: 1.25\n output_price: 10\n supports_vision: true\n supports_function_calling: true\n - name: openai/gpt-5-mini\n max_input_tokens: 400000\n max_output_tokens: 128000\n input_price: 0.25\n output_price: 2\n supports_vision: true\n supports_function_calling: true\n - name: openai/gpt-5-nano\n max_input_tokens: 400000\n max_output_tokens: 128000\n input_price: 0.05\n output_price: 0.4\n supports_vision: true\n supports_function_calling: true\n - name: openai/gpt-4.1\n max_input_tokens: 1047576\n max_output_tokens: 32768\n input_price: 2\n output_price: 8\n supports_vision: true\n supports_function_calling: true\n - name: openai/gpt-4o\n max_input_tokens: 128000\n input_price: 2.5\n output_price: 10\n supports_vision: true\n supports_function_calling: true\n - name: openai/gpt-oss-120b\n max_input_tokens: 131072\n input_price: 0.09\n output_price: 0.45\n supports_function_calling: true\n - name: openai/gpt-oss-20b\n max_input_tokens: 131072\n input_price: 0.04\n output_price: 0.16\n supports_function_calling: true\n - name: google/gemini-2.5-flash\n max_input_tokens: 1048576\n input_price: 0.3\n output_price: 2.5\n supports_vision: true\n supports_function_calling: true\n - name: google/gemini-2.5-pro\n max_input_tokens: 1048576\n input_price: 1.25\n output_price: 10\n supports_vision: true\n supports_function_calling: true\n - name: google/gemini-2.5-flash-lite\n max_input_tokens: 1048576\n input_price: 0.3\n output_price: 0.4\n supports_vision: true\n - name: google/gemini-2.0-flash-001\n max_input_tokens: 1000000\n input_price: 0.15\n output_price: 0.6\n supports_vision: true\n supports_function_calling: true\n - name: google/gemini-2.0-flash-lite-001\n max_input_tokens: 1048576\n input_price: 0.075\n output_price: 0.3\n supports_vision: true\n supports_function_calling: true\n - name: google/gemma-3-27b-it\n max_input_tokens: 131072\n input_price: 0.1\n output_price: 0.2\n - name: anthropic/claude-opus-4.6\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n supports_function_calling: true\n - name: anthropic/claude-sonnet-4.6\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n supports_function_calling: true\n - name: anthropic/claude-opus-4.5\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 5\n output_price: 25\n supports_vision: true\n supports_function_calling: true\n - name: anthropic/claude-sonnet-4.5\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 3\n output_price: 15\n supports_vision: true\n supports_function_calling: true\n - name: anthropic/claude-haiku-4.5\n max_input_tokens: 200000\n max_output_tokens: 8192\n require_max_tokens: true\n input_price: 1\n output_price: 5\n supports_vision: true\n supports_function_calling: true\n - name: meta-llama/llama-4-maverick\n max_input_tokens: 1048576\n input_price: 0.18\n output_price: 0.6\n supports_vision: true\n supports_function_calling: true\n - name: meta-llama/llama-4-scout\n max_input_tokens: 327680\n input_price: 0.08\n output_price: 0.3\n supports_vision: true\n supports_function_calling: true\n - name: meta-llama/llama-3.3-70b-instruct\n max_input_tokens: 131072\n input_price: 0.12\n output_price: 0.3\n - name: mistralai/mistral-large-2512\n max_input_tokens: 262144\n input_price: 0.5\n output_price: 1.5\n supports_function_calling: true\n - name: mistralai/mistral-medium-3.1\n max_input_tokens: 131072\n input_price: 0.4\n output_price: 2\n supports_function_calling: true\n supports_vision: true\n - name: mistralai/mistral-small-3.2-24b-instruct\n max_input_tokens: 131072\n input_price: 0.1\n output_price: 0.3\n supports_vision: true\n - name: mistralai/devstral-2512\n max_input_tokens: 262144\n input_price: 0.5\n output_price: 0.22\n supports_function_calling: true\n - name: mistralai/devstral-small\n max_input_tokens: 131072\n input_price: 0.07\n output_price: 0.28\n supports_function_calling: true\n - name: mistralai/codestral-2508\n max_input_tokens: 256000\n input_price: 0.3\n output_price: 0.9\n supports_function_calling: true\n - name: mistralai/ministral-14b-2512\n max_input_tokens: 262144\n input_price: 0.2\n output_price: 0.2\n supports_function_calling: true\n - name: ai21/jamba-large-1.7\n max_input_tokens: 256000\n input_price: 2\n output_price: 8\n supports_function_calling: true\n - name: cohere/command-a\n max_input_tokens: 256000\n input_price: 2.5\n output_price: 10\n supports_function_calling: true\n - name: cohere/command-r7b-12-2024\n max_input_tokens: 128000\n max_output_tokens: 4096\n input_price: 0.0375\n output_price: 0.15\n - name: deepseek/deepseek-v3.2\n max_input_tokens: 163840\n input_price: 0.25\n output_price: 0.38\n - name: qwen/qwen3-max\n max_input_tokens: 262144\n input_price: 1.2\n output_price: 6\n supports_function_calling: true\n - name: qwen/qwen3-max-thinking\n max_input_tokens: 262144\n input_price: 1.2\n output_price: 6\n supports_function_calling: true\n - name: qwen/qwen3.5-plus-02-15\n max_input_tokens: 1000000\n max_output_tokens: 8192\n input_price: 0.4\n output_price: 2.4\n supports_function_calling: true\n - name: qwen/qwen3.5-397b-a17b\n max_input_tokens: 262144\n max_output_tokens: 8192\n input_price: 0.15\n output_price: 1\n supports_function_calling: true\n - name: qwen/qwen3-next-80b-a3b-instruct\n max_input_tokens: 262144\n input_price: 0.1\n output_price: 0.8\n supports_function_calling: true\n - name: qwen/qwen3-next-80b-a3b-thinking\n max_input_tokens: 262144\n input_price: 0.1\n output_price: 0.8\n - name: qwen/qwen3-235b-a22b-2507 # Qwen3 235B A22B Instruct 2507\n max_input_tokens: 262144\n input_price: 0.12\n output_price: 0.59\n supports_function_calling: true\n - name: qwen/qwen3-235b-a22b-thinking-2507\n max_input_tokens: 262144\n input_price: 0.118\n output_price: 0.118\n - name: qwen/qwen3-30b-a3b-instruct-2507\n max_input_tokens: 131072\n input_price: 0.2\n output_price: 0.8\n - name: qwen/qwen3-30b-a3b-thinking-2507\n max_input_tokens: 262144\n input_price: 0.071\n output_price: 0.285\n - name: qwen/qwen3-vl-32b-instruct\n max_input_tokens: 262144\n input_price: 0.35\n output_price: 1.1\n supports_vision: true\n - name: qwen/qwen3-vl-8b-instruct\n max_input_tokens: 262144\n input_price: 0.08\n output_price: 0.50\n supports_vision: true\n - name: qwen/qwen3-coder-next\n max_input_tokens: 262144\n input_price: 0.12\n output_price: 0.75\n supports_function_calling: true\n - name: qwen/qwen3-coder-plus\n max_input_tokens: 128000\n input_price: 1\n output_price: 5\n supports_function_calling: true\n - name: qwen/qwen3-coder-flash\n max_input_tokens: 128000\n input_price: 0.3\n output_price: 1.5\n supports_function_calling: true\n - name: qwen/qwen3-coder # Qwen3 Coder 480B A35B\n max_input_tokens: 262144\n input_price: 0.22\n output_price: 0.95\n supports_function_calling: true\n - name: qwen/qwen3-coder-30b-a3b-instruct\n max_input_tokens: 262144\n input_price: 0.052\n output_price: 0.207\n supports_function_calling: true\n - name: moonshotai/kimi-k2.5\n max_input_tokens: 262144\n input_price: 0.57\n output_price: 2.85\n supports_vision: true\n supports_function_calling: true\n - name: moonshotai/kimi-k2-0905\n max_input_tokens: 262144\n input_price: 0.296\n output_price: 1.185\n supports_vision: true\n supports_function_calling: true\n - name: moonshotai/kimi-k2-thinking\n max_input_tokens: 262144\n input_price: 0.45\n output_price: 2.35\n supports_function_calling: true\n - name: x-ai/grok-4.1-fast\n max_input_tokens: 2000000\n input_price: 0.2\n output_price: 0.5\n supports_function_calling: true\n - name: x-ai/grok-code-fast-1\n max_input_tokens: 256000\n input_price: 0.2\n output_price: 1.5\n supports_function_calling: true\n - name: amazon/nova-premier-v1\n max_input_tokens: 1000000\n input_price: 2.5\n output_price: 12.5\n supports_vision: true\n - name: amazon/nova-pro-v1\n max_input_tokens: 300000\n max_output_tokens: 5120\n input_price: 0.8\n output_price: 3.2\n supports_vision: true\n - name: amazon/nova-lite-v1\n max_input_tokens: 300000\n max_output_tokens: 5120\n input_price: 0.06\n output_price: 0.24\n supports_vision: true\n - name: amazon/nova-micro-v1\n max_input_tokens: 128000\n max_output_tokens: 5120\n input_price: 0.035\n output_price: 0.14\n - name: perplexity/sonar-pro\n max_input_tokens: 200000\n input_price: 3\n output_price: 15\n - name: perplexity/sonar\n max_input_tokens: 127072\n input_price: 1\n output_price: 1\n - name: perplexity/sonar-reasoning-pro\n max_input_tokens: 128000\n input_price: 2\n output_price: 8\n patch:\n body:\n include_reasoning: true\n - name: perplexity/sonar-deep-research\n max_input_tokens: 200000\n input_price: 2\n output_price: 8\n patch:\n body:\n include_reasoning: true\n - name: minimax/minimax-m2.5\n max_input_tokens: 196608\n input_price: 0.3\n output_price: 1.1\n supports_function_calling: true\n - name: minimax/minimax-m2.1\n max_input_tokens: 196608\n input_price: 0.12\n output_price: 0.48\n supports_function_calling: true\n - name: z-ai/glm-5\n max_input_tokens: 204800\n input_price: 0.95\n output_price: 2.55\n supports_function_calling: true\n - name: z-ai/glm-4.7\n max_input_tokens: 202752\n input_price: 0.16\n output_price: 0.80\n supports_function_calling: true\n - name: z-ai/glm-4.7-flash\n max_input_tokens: 202752\n input_price: 0.07\n output_price: 0.40\n supports_function_calling: true\n - name: z-ai/glm-4.6v\n max_input_tokens: 131072\n input_price: 0.3\n output_price: 0.9\n supports_vision: true\n\n# Links:\n# - https://github.com/marketplace?type=models\n- provider: github\n models:\n - name: gpt-5\n max_input_tokens: 400000\n max_output_tokens: 128000\n supports_vision: true\n supports_function_calling: true\n - name: gpt-5-mini\n max_input_tokens: 400000\n max_output_tokens: 128000\n supports_vision: true\n supports_function_calling: true\n - name: gpt-5-nano\n max_input_tokens: 400000\n max_output_tokens: 128000\n supports_vision: true\n supports_function_calling: true\n - name: gpt-4.1\n max_input_tokens: 1047576\n max_output_tokens: 32768\n supports_vision: true\n supports_function_calling: true\n - name: gpt-4o\n max_input_tokens: 128000\n max_output_tokens: 16384\n supports_function_calling: true\n - name: text-embedding-3-large\n type: embedding\n max_tokens_per_chunk: 8191\n default_chunk_size: 2000\n max_batch_size: 100\n - name: text-embedding-3-small\n type: embedding\n max_tokens_per_chunk: 8191\n default_chunk_size: 2000\n max_batch_size: 100\n - name: llama-4-maverick-17b-128e-instruct-fp8\n max_input_tokens: 1048576\n supports_vision: true\n - name: llama-4-scout-17b-16e-instruct\n max_input_tokens: 327680\n supports_vision: true\n - name: llama-3.3-70b-instruct\n max_input_tokens: 131072\n - name: mistral-medium-2505\n max_input_tokens: 131072\n supports_function_calling: true\n - name: mistral-small-2503\n max_input_tokens: 131072\n supports_function_calling: true\n - name: codestral-2501\n max_input_tokens: 256000\n supports_function_calling: true\n - name: cohere-embed-v3-english\n type: embedding\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 96\n - name: cohere-embed-v3-multilingual\n type: embedding\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 96\n - name: deepseek-r1-0528\n max_input_tokens: 163840\n - name: deepseek-v3-0324\n max_input_tokens: 163840\n - name: mai-ds-r1\n max_input_tokens: 163840\n - name: phi-4\n max_input_tokens: 16384\n - name: phi-4-mini-instruct\n max_input_tokens: 131072\n - name: phi-4-reasoning\n max_input_tokens: 33792\n - name: phi-4-mini-reasoning\n max_input_tokens: 131072\n - name: grok-3\n max_input_tokens: 131072\n - name: grok-3-mini\n max_input_tokens: 131072\n\n# Links:\n# - https://deepinfra.com/models\n# - https://deepinfra.com/docs/openai_api\n- provider: deepinfra\n models:\n - name: openai/gpt-oss-120b\n max_input_tokens: 131072\n input_price: 0.09\n output_price: 0.45\n supports_function_calling: true\n - name: openai/gpt-oss-20b\n max_input_tokens: 131072\n input_price: 0.04\n output_price: 0.16\n supports_function_calling: true\n - name: meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8\n max_input_tokens: 1048576\n input_price: 0.18\n output_price: 0.6\n supports_vision: true\n - name: meta-llama/Llama-4-Scout-17B-16E-Instruct\n max_input_tokens: 327680\n input_price: 0.08\n output_price: 0.3\n supports_vision: true\n - name: Qwen/Qwen3-Max\n max_input_tokens: 262144\n input_price: 1.2\n output_price: 6\n supports_function_calling: true\n - name: Qwen/Qwen3-Max-Thinking\n max_input_tokens: 262144\n input_price: 1.2\n output_price: 6\n supports_function_calling: true\n - name: Qwen/Qwen3-Next-80B-A3B-Instruct\n max_input_tokens: 262144\n input_price: 0.14\n output_price: 1.4\n supports_function_calling: true\n - name: Qwen/Qwen3-Next-80B-A3B-Thinking\n max_input_tokens: 262144\n input_price: 0.14\n output_price: 1.4\n - name: Qwen/Qwen3-235B-A22B-Instruct-2507\n max_input_tokens: 131072\n input_price: 0.13\n output_price: 0.6\n supports_function_calling: true\n - name: Qwen/Qwen3-235B-A22B-Thinking-2507\n max_input_tokens: 131072\n input_price: 0.13\n output_price: 0.6\n - name: Qwen/Qwen3-Coder-480B-A35B-Instruct\n max_input_tokens: 131072\n input_price: 0.4\n output_price: 1.6\n supports_function_calling: true\n - name: Qwen/Qwen3-Coder-30B-A3B-Instruct\n max_input_tokens: 262144\n input_price: 0.07\n output_price: 0.27\n supports_function_calling: true\n - name: Qwen/Qwen3-30B-A3B\n max_input_tokens: 40960\n input_price: 0.1\n output_price: 0.3\n - name: Qwen/Qwen3-VL-8B-Instruct\n max_input_tokens: 262144\n input_price: 0.18\n output_price: 0.69\n supports_vision: true\n - name: deepseek-ai/DeepSeek-V3.2\n max_input_tokens: 163840\n input_price: 0.26\n output_price: 0.39\n supports_function_calling: true\n - name: google/gemma-3-27b-it\n max_input_tokens: 131072\n input_price: 0.1\n output_price: 0.2\n - name: mistralai/Mistral-Small-3.2-24B-Instruct-2506\n max_input_tokens: 32768\n input_price: 0.06\n output_price: 0.12\n - name: moonshotai/Kimi-K2.5\n max_input_tokens: 262144\n input_price: 0.5\n output_price: 2.8\n supports_function_calling: true\n - name: moonshotai/Kimi-K2-Instruct-0905\n max_input_tokens: 262144\n input_price: 0.5\n output_price: 2.0\n supports_function_calling: true\n - name: moonshotai/Kimi-K2-Thinking\n max_input_tokens: 262144\n input_price: 0.55\n output_price: 2.5\n supports_function_calling: true\n - name: MiniMaxAI/MiniMax-M2.5\n max_input_tokens: 196608\n input_price: 0.27\n output_price: 0.95\n supports_function_calling: true\n - name: MiniMaxAI/MiniMax-M2.1\n max_input_tokens: 196608\n input_price: 0.27\n output_price: 0.95\n supports_function_calling: true\n - name: zai-org/GLM-5\n max_input_tokens: 202752\n input_price: 0.8\n output_price: 2.56\n supports_function_calling: true\n - name: zai-org/GLM-4.7\n max_input_tokens: 202752\n input_price: 0.43\n output_price: 1.75\n supports_function_calling: true\n - name: zai-org/GLM-4.7-Flash\n max_input_tokens: 202752\n input_price: 0.06\n output_price: 0.4\n supports_function_calling: true\n - name: zai-org/GLM-4.6V\n max_input_tokens: 131072\n input_price: 0.3\n output_price: 0.9\n supports_vision: true\n - name: BAAI/bge-large-en-v1.5\n type: embedding\n input_price: 0.01\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 100\n - name: BAAI/bge-m3\n type: embedding\n input_price: 0.01\n max_tokens_per_chunk: 8192\n default_chunk_size: 2000\n max_batch_size: 100\n - name: intfloat/e5-large-v2\n type: embedding\n input_price: 0.01\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 100\n - name: intfloat/multilingual-e5-large\n type: embedding\n input_price: 0.01\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 100\n - name: thenlper/gte-large\n type: embedding\n input_price: 0.01\n max_tokens_per_chunk: 512\n default_chunk_size: 1000\n max_batch_size: 100\n\n# Links:\n# - https://jina.ai/models\n# - https://api.jina.ai/redoc\n- provider: jina\n models:\n - name: jina-embeddings-v3\n type: embedding\n input_price: 0\n max_tokens_per_chunk: 8192\n default_chunk_size: 2000\n max_batch_size: 100\n - name: jina-clip-v2\n type: embedding\n input_price: 0\n max_tokens_per_chunk: 8192\n default_chunk_size: 1500\n max_batch_size: 100\n - name: jina-colbert-v2\n type: embedding\n input_price: 0\n max_tokens_per_chunk: 8192\n default_chunk_size: 1500\n max_batch_size: 100\n - name: jina-reranker-v2-base-multilingual\n type: reranker\n max_input_tokens: 8192\n input_price: 0\n - name: jina-colbert-v2\n type: reranker\n max_input_tokens: 8192\n input_price: 0\n\n# Links:\n# - https://docs.voyageai.com/docs/embeddings\n# - https://docs.voyageai.com/docs/pricing\n# - https://docs.voyageai.com/reference/\n- provider: voyageai\n models:\n - name: voyage-3-large\n type: embedding\n max_input_tokens: 120000\n input_price: 0.18\n max_tokens_per_chunk: 32000\n default_chunk_size: 2000\n max_batch_size: 128\n - name: voyage-3\n type: embedding\n max_input_tokens: 320000\n input_price: 0.06\n max_tokens_per_chunk: 32000\n default_chunk_size: 2000\n max_batch_size: 128\n - name: voyage-3-lite\n type: embedding\n max_input_tokens: 1000000\n input_price: 0.02\n max_tokens_per_chunk: 32000\n default_chunk_size: 1000\n max_batch_size: 128\n - name: rerank-2\n type: reranker\n max_input_tokens: 16000\n input_price: 0.05\n - name: rerank-2-lite\n type: reranker\n max_input_tokens: 8000\n input_price: 0.02\n" }, { "path": "scripts/completions/aichat.bash", "content": "_aichat() {\n local cur prev words cword i opts cmd\n COMPREPLY=()\n\n _get_comp_words_by_ref -n : cur prev words cword\n\n for i in ${words[@]}\n do\n case \"${cmd},${i}\" in\n \",$1\")\n cmd=\"aichat\"\n ;;\n *)\n ;;\n esac\n done\n\n case \"${cmd}\" in\n aichat)\n opts=\"-m -r -s -a -e -c -f -S -h -V --model --prompt --role --session --empty-session --save-session --agent --agent-variable --rag --rebuild-rag --macro --serve --execute --code --file --no-stream --dry-run --info --sync-models --list-models --list-roles --list-sessions --list-agents --list-rags --list-macros --help --version\"\n if [[ ${cur} == -* || ${cword} -eq 1 ]] ; then\n COMPREPLY=( $(compgen -W \"${opts}\" -- \"${cur}\") )\n return 0\n fi\n\n case \"${prev}\" in\n -m|--model)\n COMPREPLY=($(compgen -W \"$(\"$1\" --list-models)\" -- \"${cur}\"))\n __ltrim_colon_completions \"$cur\"\n return 0\n ;;\n --prompt)\n COMPREPLY=($(compgen -f \"${cur}\"))\n return 0\n ;;\n -r|--role)\n COMPREPLY=($(compgen -W \"$(\"$1\" --list-roles)\" -- \"${cur}\"))\n __ltrim_colon_completions \"$cur\"\n return 0\n ;;\n -s|--session)\n COMPREPLY=($(compgen -W \"$(\"$1\" --list-sessions)\" -- \"${cur}\"))\n __ltrim_colon_completions \"$cur\"\n return 0\n ;;\n -a|--agent)\n COMPREPLY=($(compgen -W \"$(\"$1\" --list-agents)\" -- \"${cur}\"))\n __ltrim_colon_completions \"$cur\"\n return 0\n ;;\n -R|--rag)\n COMPREPLY=($(compgen -W \"$(\"$1\" --list-rags)\" -- \"${cur}\"))\n __ltrim_colon_completions \"$cur\"\n return 0\n ;;\n --macro)\n COMPREPLY=($(compgen -W \"$(\"$1\" --list-macros)\" -- \"${cur}\"))\n __ltrim_colon_completions \"$cur\"\n return 0\n ;;\n -f|--file)\n local oldifs\n if [[ -v IFS ]]; then\n oldifs=\"$IFS\"\n fi\n IFS=$'\\n'\n COMPREPLY=($(compgen -f \"${cur}\"))\n if [[ -v oldifs ]]; then\n IFS=\"$oldifs\"\n fi\n if [[ \"${BASH_VERSINFO[0]}\" -ge 4 ]]; then\n compopt -o filenames\n fi\n return 0\n ;;\n *)\n COMPREPLY=()\n ;;\n esac\n COMPREPLY=( $(compgen -W \"${opts}\" -- \"${cur}\") )\n return 0\n ;;\n esac\n}\n\nif [[ \"${BASH_VERSINFO[0]}\" -eq 4 && \"${BASH_VERSINFO[1]}\" -ge 4 || \"${BASH_VERSINFO[0]}\" -gt 4 ]]; then\n complete -F _aichat -o nosort -o bashdefault -o default aichat\nelse\n complete -F _aichat -o bashdefault -o default aichat\nfi\n" }, { "path": "scripts/completions/aichat.fish", "content": "complete -c aichat -s m -l model -x -a \"(aichat --list-models)\" -d 'Select a LLM model' -r\ncomplete -c aichat -l prompt -d 'Use the system prompt'\ncomplete -c aichat -s r -l role -x -a \"(aichat --list-roles)\" -d 'Select a role' -r\ncomplete -c aichat -s s -l session -x -a \"(aichat --list-sessions)\" -d 'Start or join a session' -r\ncomplete -c aichat -l empty-session -d 'Ensure the session is empty'\ncomplete -c aichat -l save-session -d 'Ensure the new conversation is saved to the session'\ncomplete -c aichat -s a -l agent -x -a \"(aichat --list-agents)\" -d 'Start a agent' -r\ncomplete -c aichat -l agent-variable -d 'Set agent variables'\ncomplete -c aichat -l rag -x -a\"(aichat --list-rags)\" -d 'Start a RAG' -r\ncomplete -c aichat -l rebuild-rag -d 'Rebuild the RAG to sync document changes'\ncomplete -c aichat -l macro -x -a\"(aichat --list-macros)\" -d 'Execute a macro' -r\ncomplete -c aichat -l serve -d 'Serve the LLM API and WebAPP'\ncomplete -c aichat -s e -l execute -d 'Execute commands in natural language'\ncomplete -c aichat -s c -l code -d 'Output code only'\ncomplete -c aichat -s f -l file -d 'Include files, directories, or URLs' -r -F\ncomplete -c aichat -s S -l no-stream -d 'Turn off stream mode'\ncomplete -c aichat -l dry-run -d 'Display the message without sending it'\ncomplete -c aichat -l info -d 'Display information'\ncomplete -c aichat -l sync-models -d 'Sync models updates'\ncomplete -c aichat -l list-models -d 'List all available chat models'\ncomplete -c aichat -l list-roles -d 'List all roles'\ncomplete -c aichat -l list-sessions -d 'List all sessions'\ncomplete -c aichat -l list-agents -d 'List all agents'\ncomplete -c aichat -l list-rags -d 'List all RAGs'\ncomplete -c aichat -l list-macros -d 'List all macros'\ncomplete -c aichat -s h -l help -d 'Print help'\ncomplete -c aichat -s V -l version -d 'Print version'\n" }, { "path": "scripts/completions/aichat.nu", "content": "module completions {\n\n def \"nu-complete aichat completions\" [] {\n [ \"bash\" \"zsh\" \"fish\" \"powershell\" \"nushell\" ]\n }\n\n def \"nu-complete aichat model\" [] {\n ^aichat --list-models |\n | lines \n | parse \"{value}\" \n }\n\n def \"nu-complete aichat role\" [] {\n ^aichat --list-roles |\n | lines \n | parse \"{value}\" \n }\n\n def \"nu-complete aichat session\" [] {\n ^aichat --list-sessions |\n | lines \n | parse \"{value}\" \n }\n\n def \"nu-complete aichat agent\" [] {\n ^aichat --list-agents |\n | lines \n | parse \"{value}\" \n }\n\n def \"nu-complete aichat rag\" [] {\n ^aichat --list-rags |\n | lines \n | parse \"{value}\" \n }\n\n def \"nu-complete aichat macro\" [] {\n ^aichat --list-macros |\n | lines \n | parse \"{value}\" \n }\n\n export extern aichat [\n --model(-m): string@\"nu-complete aichat model\" # Select a LLM model\n --prompt # Use the system prompt\n --role(-r): string@\"nu-complete aichat role\" # Select a role\n --session(-s): string@\"nu-complete aichat session\" # Start or join a session\n --empty-session # Ensure the session is empty\n --save-session # Ensure the new conversation is saved to the session\n --agent(-a): string@\"nu-complete aichat agent\" # Start a agent\n --agent-variable # Set agent variables\n --rag: string@\"nu-complete aichat rag\" # Start a RAG\n --rebuild-rag # Rebuild the RAG to sync document changes\n --macro: string@\"nu-complete aichat macro\" # Execute a macro\n --serve # Serve the LLM API and WebAPP\n --execute(-e) # Execute commands in natural language\n --code(-c) # Output code only\n --file(-f): string # Include files, directories, or URLs\n --no-stream(-S) # Turn off stream mode\n --dry-run # Display the message without sending it\n --info # Display information\n --sync-models # Sync models updates\n --list-models # List all available chat models\n --list-roles # List all roles\n --list-sessions # List all sessions\n --list-agents # List all agents\n --list-rags # List all RAGs\n --list-macros # List all macros\n ...text: string # Input text\n --help(-h) # Print help\n --version(-V) # Print version\n ]\n\n}\n\nexport use completions *\n" }, { "path": "scripts/completions/aichat.ps1", "content": "using namespace System.Management.Automation\nusing namespace System.Management.Automation.Language\n\nRegister-ArgumentCompleter -Native -CommandName 'aichat' -ScriptBlock {\n param($wordToComplete, $commandAst, $cursorPosition)\n\n $commandElements = $commandAst.CommandElements\n $command = @(\n 'aichat'\n for ($i = 1; $i -lt $commandElements.Count; $i++) {\n $element = $commandElements[$i]\n if ($element -isnot [StringConstantExpressionAst] -or\n $element.StringConstantType -ne [StringConstantType]::BareWord -or\n $element.Value.StartsWith('-') -or\n $element.Value -eq $wordToComplete) {\n break\n }\n $element.Value\n }) -join ';'\n\n $completions = @(switch ($command) {\n 'aichat' {\n [CompletionResult]::new('-m', '-m', [CompletionResultType]::ParameterName, 'Select a LLM model')\n [CompletionResult]::new('--model', '--model', [CompletionResultType]::ParameterName, 'Select a LLM model')\n [CompletionResult]::new('--prompt', '--prompt', [CompletionResultType]::ParameterName, 'Use the system prompt')\n [CompletionResult]::new('-r', '-r', [CompletionResultType]::ParameterName, 'Select a role')\n [CompletionResult]::new('--role', '--role', [CompletionResultType]::ParameterName, 'Select a role')\n [CompletionResult]::new('-s', '-s', [CompletionResultType]::ParameterName, 'Start or join a session')\n [CompletionResult]::new('--session', '--session', [CompletionResultType]::ParameterName, 'Start or join a session')\n [CompletionResult]::new('--empty-session', '--empty-session', [CompletionResultType]::ParameterName, 'Ensure the session is empty')\n [CompletionResult]::new('--save-session', '--save-session', [CompletionResultType]::ParameterName, 'Ensure the new conversation is saved to the session')\n [CompletionResult]::new('-a', '-a', [CompletionResultType]::ParameterName, 'Start a agent')\n [CompletionResult]::new('--agent', '--agent', [CompletionResultType]::ParameterName, 'Start a agent')\n [CompletionResult]::new('--agent-variable', '--agent-variable', [CompletionResultType]::ParameterName, 'Set agent variables')\n [CompletionResult]::new('--rag', '--rag', [CompletionResultType]::ParameterName, 'Start a RAG')\n [CompletionResult]::new('--rebuild-rag', '--rebuild-rag', [CompletionResultType]::ParameterName, 'Rebuild the RAG to sync document changes')\n [CompletionResult]::new('--macro', '--macro', [CompletionResultType]::ParameterName, 'Execute a macro')\n [CompletionResult]::new('--serve', '--serve', [CompletionResultType]::ParameterName, 'Serve the LLM API and WebAPP')\n [CompletionResult]::new('-e', '-e', [CompletionResultType]::ParameterName, 'Execute commands in natural language')\n [CompletionResult]::new('--execute', '--execute', [CompletionResultType]::ParameterName, 'Execute commands in natural language')\n [CompletionResult]::new('-c', '-c', [CompletionResultType]::ParameterName, 'Output code only')\n [CompletionResult]::new('--code', '--code', [CompletionResultType]::ParameterName, 'Output code only')\n [CompletionResult]::new('-f', '-f', [CompletionResultType]::ParameterName, 'Include files, directories, or URLs')\n [CompletionResult]::new('--file', '--file', [CompletionResultType]::ParameterName, 'Include files, directories, or URLs')\n [CompletionResult]::new('-S', '-S', [CompletionResultType]::ParameterName, 'Turn off stream mode')\n [CompletionResult]::new('--no-stream', '--no-stream', [CompletionResultType]::ParameterName, 'Turn off stream mode')\n [CompletionResult]::new('--dry-run', '--dry-run', [CompletionResultType]::ParameterName, 'Display the message without sending it')\n [CompletionResult]::new('--info', '--info', [CompletionResultType]::ParameterName, 'Display information')\n [CompletionResult]::new('--sync-models', '--sync-models', [CompletionResultType]::ParameterName, 'Sync models updates')\n [CompletionResult]::new('--list-models', '--list-models', [CompletionResultType]::ParameterName, 'List all available chat models')\n [CompletionResult]::new('--list-roles', '--list-roles', [CompletionResultType]::ParameterName, 'List all roles')\n [CompletionResult]::new('--list-sessions', '--list-sessions', [CompletionResultType]::ParameterName, 'List all sessions')\n [CompletionResult]::new('--list-agents', '--list-agents', [CompletionResultType]::ParameterName, 'List all agents')\n [CompletionResult]::new('--list-rags', '--list-rags', [CompletionResultType]::ParameterName, 'List all RAGs')\n [CompletionResult]::new('--list-macros', '--list-macros', [CompletionResultType]::ParameterName, 'List all macros')\n [CompletionResult]::new('-h', '-h', [CompletionResultType]::ParameterName, 'Print help')\n [CompletionResult]::new('--help', '--help', [CompletionResultType]::ParameterName, 'Print help')\n [CompletionResult]::new('-V', '-V', [CompletionResultType]::ParameterName, 'Print version')\n [CompletionResult]::new('--version', '--version', [CompletionResultType]::ParameterName, 'Print version')\n break\n }\n })\n\n function Get-AichatValues($arg) {\n $(aichat $arg) -split '\\n' | ForEach-Object { [CompletionResult]::new($_) }\n }\n\n if ($commandElements.Count -gt 1) {\n $offset=2\n if ($wordToComplete -eq \"\") {\n $offset=1\n }\n $flag = $commandElements[$commandElements.Count-$offset].ToString()\n dump-args $flag ($flag -eq \"-R\") > /tmp/file1\n if ($flag -ceq \"-m\" -or $flag -eq \"--model\") {\n $completions = Get-AichatValues \"--list-models\"\n } elseif ($flag -ceq \"-r\" -or $flag -eq \"--role\") {\n $completions = Get-AichatValues \"--list-roles\"\n } elseif ($flag -ceq \"-s\" -or $flag -eq \"--session\") {\n $completions = Get-AichatValues \"--list-sessions\"\n } elseif ($flag -ceq \"-a\" -or $flag -eq \"--agent\") {\n $completions = Get-AichatValues \"--list-agents\"\n } elseif ($flag -eq \"--rag\") {\n $completions = Get-AichatValues \"--list-rags\"\n } elseif ($flag -eq \"--macro\") {\n $completions = Get-AichatValues \"--list-macros\"\n } elseif ($flag -ceq \"-f\" -or $flag -eq \"--file\") {\n $completions = @()\n }\n }\n\n $completions.Where{ $_.CompletionText -like \"$wordToComplete*\" } |\n Sort-Object -Property ListItemText\n}\n" }, { "path": "scripts/completions/aichat.zsh", "content": "#compdef aichat\n\nautoload -U is-at-least\n\n_aichat() {\n typeset -A opt_args\n typeset -a _arguments_options\n local ret=1\n\n if is-at-least 5.2; then\n _arguments_options=(-s -S -C)\n else\n _arguments_options=(-s -C)\n fi\n\n local context curcontext=\"$curcontext\" state line\n local common=(\n'-m[Select a LLM model]:MODEL:->models' \\\n'--model[Select a LLM model]:MODEL:->models' \\\n'--prompt[Use the system prompt]:PROMPT: ' \\\n'-r[Select a role]:ROLE:->roles' \\\n'--role[Select a role]:ROLE:->roles' \\\n'-s[Start or join a session]:SESSION:->sessions' \\\n'--session[Start or join a session]:SESSION:->sessions' \\\n'--empty-session[Ensure the session is empty]' \\\n'--save-session[Ensure the new conversation is saved to the session]' \\\n'-a[Start a agent]:AGENT:->agents' \\\n'--agent[Start a agent]:AGENT:->agents' \\\n'--agent-variable[Set agent variables]' \\\n'--rag[Start a RAG]:RAG:->rags' \\\n'--rebuild-rag[Rebuild the RAG to sync document changes]' \\\n'--macro[Execute a macro]:MACRO:->macros' \\\n'--serve[Serve the LLM API and WebAPP]' \\\n'-e[Execute commands in natural language]' \\\n'--execute[Execute commands in natural language]' \\\n'-c[Output code only]' \\\n'--code[Output code only]' \\\n'*-f[Include files, directories, or URLs]:FILE:_files' \\\n'*--file[Include files, directories, or URLs]:FILE:_files' \\\n'-S[Turn off stream mode]' \\\n'--no-stream[Turn off stream mode]' \\\n'--dry-run[Display the message without sending it]' \\\n'--info[Display information]' \\\n'--sync-models[Sync models updates]' \\\n'--list-models[List all available chat models]' \\\n'--list-roles[List all roles]' \\\n'--list-sessions[List all sessions]' \\\n'--list-agents[List all agents]' \\\n'--list-rags[List all RAGs]' \\\n'--list-macros[List all macros]' \\\n'-h[Print help]' \\\n'--help[Print help]' \\\n'-V[Print version]' \\\n'--version[Print version]' \\\n'*::text -- Input text:' \\\n )\n\n\n _arguments \"${_arguments_options[@]}\" $common \\\n && ret=0 \n case $state in\n models|roles|sessions|agents|rags|macros)\n local -a values expl\n values=( ${(f)\"$(_call_program values aichat --list-$state)\"} )\n _wanted values expl $state compadd -a values && ret=0\n ;;\n esac\n return ret\n}\n\n(( $+functions[_aichat_commands] )) ||\n_aichat_commands() {\n local commands; commands=()\n _describe -t commands 'aichat commands' commands \"$@\"\n}\n\nif [ \"$funcstack[1]\" = \"_aichat\" ]; then\n _aichat \"$@\"\nelse\n compdef _aichat aichat\nfi\n" }, { "path": "scripts/shell-integration/integration.bash", "content": "_aichat_bash() {\n if [[ -n \"$READLINE_LINE\" ]]; then\n READLINE_LINE=$(aichat -e \"$READLINE_LINE\")\n READLINE_POINT=${#READLINE_LINE}\n fi\n}\nbind -x '\"\\ee\": _aichat_bash'" }, { "path": "scripts/shell-integration/integration.fish", "content": "function _aichat_fish\n set -l _old (commandline)\n if test -n $_old\n echo -n \"⌛\"\n commandline -f repaint\n commandline (aichat -e $_old)\n end\nend\nbind \\ee _aichat_fish" }, { "path": "scripts/shell-integration/integration.nu", "content": "def _aichat_nushell [] {\n let _prev = (commandline)\n if ($_prev != \"\") {\n print '⌛'\n commandline edit -r (aichat -e $_prev)\n }\n}\n\n$env.config.keybindings = ($env.config.keybindings | append {\n name: aichat_integration\n modifier: alt\n keycode: char_e\n mode: [emacs, vi_insert]\n event:[\n {\n send: executehostcommand,\n cmd: \"_aichat_nushell\"\n }\n ]\n }\n)" }, { "path": "scripts/shell-integration/integration.ps1", "content": "Set-PSReadLineKeyHandler -Chord \"alt+e\" -ScriptBlock {\n $_old = $null\n [Microsoft.PowerShell.PSConsoleReadline]::GetBufferState([ref]$_old, [ref]$null)\n if ($_old) {\n [Microsoft.PowerShell.PSConsoleReadLine]::Insert('⌛')\n $_new = (aichat -e $_old)\n [Microsoft.PowerShell.PSConsoleReadLine]::DeleteLine()\n [Microsoft.PowerShell.PSConsoleReadline]::Insert($_new)\n }\n}" }, { "path": "scripts/shell-integration/integration.zsh", "content": "_aichat_zsh() {\n if [[ -n \"$BUFFER\" ]]; then\n local _old=$BUFFER\n BUFFER+=\"⌛\"\n zle -I && zle redisplay\n BUFFER=$(aichat -e \"$_old\")\n zle end-of-line\n fi\n}\nzle -N _aichat_zsh\nbindkey '\\ee' _aichat_zsh" }, { "path": "src/cli.rs", "content": "use anyhow::{Context, Result};\nuse clap::Parser;\nuse is_terminal::IsTerminal;\nuse std::io::{stdin, Read};\n\n#[derive(Parser, Debug)]\n#[command(author, version, about, long_about = None)]\npub struct Cli {\n /// Select a LLM model\n #[clap(short, long)]\n pub model: Option,\n /// Use the system prompt\n #[clap(long)]\n pub prompt: Option,\n /// Select a role\n #[clap(short, long)]\n pub role: Option,\n /// Start or join a session\n #[clap(short = 's', long)]\n pub session: Option>,\n /// Ensure the session is empty\n #[clap(long)]\n pub empty_session: bool,\n /// Ensure the new conversation is saved to the session\n #[clap(long)]\n pub save_session: bool,\n /// Start a agent\n #[clap(short = 'a', long)]\n pub agent: Option,\n /// Set agent variables\n #[clap(long, value_names = [\"NAME\", \"VALUE\"], num_args = 2)]\n pub agent_variable: Vec,\n /// Start a RAG\n #[clap(long)]\n pub rag: Option,\n /// Rebuild the RAG to sync document changes\n #[clap(long)]\n pub rebuild_rag: bool,\n /// Execute a macro\n #[clap(long = \"macro\", value_name = \"MACRO\")]\n pub macro_name: Option,\n /// Serve the LLM API and WebAPP\n #[clap(long, value_name = \"ADDRESS\")]\n pub serve: Option>,\n /// Execute commands in natural language\n #[clap(short = 'e', long)]\n pub execute: bool,\n /// Output code only\n #[clap(short = 'c', long)]\n pub code: bool,\n /// Include files, directories, or URLs\n #[clap(short = 'f', long, value_name = \"FILE\")]\n pub file: Vec,\n /// Turn off stream mode\n #[clap(short = 'S', long)]\n pub no_stream: bool,\n /// Display the message without sending it\n #[clap(long)]\n pub dry_run: bool,\n /// Display information\n #[clap(long)]\n pub info: bool,\n /// Sync models updates\n #[clap(long)]\n pub sync_models: bool,\n /// List all available chat models\n #[clap(long)]\n pub list_models: bool,\n /// List all roles\n #[clap(long)]\n pub list_roles: bool,\n /// List all sessions\n #[clap(long)]\n pub list_sessions: bool,\n /// List all agents\n #[clap(long)]\n pub list_agents: bool,\n /// List all RAGs\n #[clap(long)]\n pub list_rags: bool,\n /// List all macros\n #[clap(long)]\n pub list_macros: bool,\n /// Input text\n #[clap(trailing_var_arg = true)]\n text: Vec,\n}\n\nimpl Cli {\n pub fn text(&self) -> Result> {\n let mut stdin_text = String::new();\n if !stdin().is_terminal() {\n let _ = stdin()\n .read_to_string(&mut stdin_text)\n .context(\"Invalid stdin pipe\")?;\n };\n match self.text.is_empty() {\n true => {\n if stdin_text.is_empty() {\n Ok(None)\n } else {\n Ok(Some(stdin_text))\n }\n }\n false => {\n if self.macro_name.is_some() {\n let text = self\n .text\n .iter()\n .map(|v| shell_words::quote(v))\n .collect::>()\n .join(\" \");\n if stdin_text.is_empty() {\n Ok(Some(text))\n } else {\n Ok(Some(format!(\"{text} -- {stdin_text}\")))\n }\n } else {\n let text = self.text.join(\" \");\n if stdin_text.is_empty() {\n Ok(Some(text))\n } else {\n Ok(Some(format!(\"{text}\\n{stdin_text}\")))\n }\n }\n }\n }\n }\n}\n" }, { "path": "src/client/access_token.rs", "content": "use anyhow::{anyhow, Result};\nuse chrono::Utc;\nuse indexmap::IndexMap;\nuse parking_lot::RwLock;\nuse std::sync::LazyLock;\n\nstatic ACCESS_TOKENS: LazyLock>> =\n LazyLock::new(|| RwLock::new(IndexMap::new()));\n\npub fn get_access_token(client_name: &str) -> Result {\n ACCESS_TOKENS\n .read()\n .get(client_name)\n .map(|(token, _)| token.clone())\n .ok_or_else(|| anyhow!(\"Invalid access token\"))\n}\n\npub fn is_valid_access_token(client_name: &str) -> bool {\n let access_tokens = ACCESS_TOKENS.read();\n let (token, expires_at) = match access_tokens.get(client_name) {\n Some(v) => v,\n None => return false,\n };\n !token.is_empty() && Utc::now().timestamp() < *expires_at\n}\n\npub fn set_access_token(client_name: &str, token: String, expires_at: i64) {\n let mut access_tokens = ACCESS_TOKENS.write();\n let entry = access_tokens.entry(client_name.to_string()).or_default();\n entry.0 = token;\n entry.1 = expires_at;\n}\n" }, { "path": "src/client/azure_openai.rs", "content": "use super::openai::*;\nuse super::*;\n\nuse anyhow::Result;\nuse serde::Deserialize;\n\n#[derive(Debug, Clone, Deserialize)]\npub struct AzureOpenAIConfig {\n pub name: Option,\n pub api_base: Option,\n pub api_key: Option,\n #[serde(default)]\n pub models: Vec,\n pub patch: Option,\n pub extra: Option,\n}\n\nimpl AzureOpenAIClient {\n config_get_fn!(api_base, get_api_base);\n config_get_fn!(api_key, get_api_key);\n\n pub const PROMPTS: [PromptAction<'static>; 2] = [\n (\n \"api_base\",\n \"API Base\",\n Some(\"e.g. https://{RESOURCE}.openai.azure.com\"),\n ),\n (\"api_key\", \"API Key\", None),\n ];\n}\n\nimpl_client_trait!(\n AzureOpenAIClient,\n (\n prepare_chat_completions,\n openai_chat_completions,\n openai_chat_completions_streaming\n ),\n (prepare_embeddings, openai_embeddings),\n (noop_prepare_rerank, noop_rerank),\n);\n\nfn prepare_chat_completions(\n self_: &AzureOpenAIClient,\n data: ChatCompletionsData,\n) -> Result {\n let api_base = self_.get_api_base()?;\n let api_key = self_.get_api_key()?;\n\n let url = format!(\n \"{}/openai/deployments/{}/chat/completions?api-version=2024-12-01-preview\",\n &api_base,\n self_.model.real_name()\n );\n\n let body = openai_build_chat_completions_body(data, &self_.model);\n\n let mut request_data = RequestData::new(url, body);\n\n request_data.header(\"api-key\", api_key);\n\n Ok(request_data)\n}\n\nfn prepare_embeddings(self_: &AzureOpenAIClient, data: &EmbeddingsData) -> Result {\n let api_base = self_.get_api_base()?;\n let api_key = self_.get_api_key()?;\n\n let url = format!(\n \"{}/openai/deployments/{}/embeddings?api-version=2024-10-21\",\n &api_base,\n self_.model.real_name()\n );\n\n let body = openai_build_embeddings_body(data, &self_.model);\n\n let mut request_data = RequestData::new(url, body);\n\n request_data.header(\"api-key\", api_key);\n\n Ok(request_data)\n}\n" }, { "path": "src/client/bedrock.rs", "content": "use super::*;\n\nuse crate::utils::{base64_decode, encode_uri, hex_encode, hmac_sha256, sha256, strip_think_tag};\n\nuse anyhow::{bail, Context, Result};\nuse aws_smithy_eventstream::frame::{DecodedFrame, MessageFrameDecoder};\nuse aws_smithy_eventstream::smithy::parse_response_headers;\nuse bytes::BytesMut;\nuse chrono::{DateTime, Utc};\nuse futures_util::StreamExt;\nuse indexmap::IndexMap;\nuse reqwest::{Client as ReqwestClient, Method, RequestBuilder};\nuse serde::Deserialize;\nuse serde_json::{json, Value};\n\n#[derive(Debug, Clone, Deserialize)]\npub struct BedrockConfig {\n pub name: Option,\n pub access_key_id: Option,\n pub secret_access_key: Option,\n pub region: Option,\n pub session_token: Option,\n #[serde(default)]\n pub models: Vec,\n pub patch: Option,\n pub extra: Option,\n}\n\nimpl BedrockClient {\n config_get_fn!(access_key_id, get_access_key_id);\n config_get_fn!(secret_access_key, get_secret_access_key);\n config_get_fn!(region, get_region);\n config_get_fn!(session_token, get_session_token);\n\n pub const PROMPTS: [PromptAction<'static>; 3] = [\n (\"access_key_id\", \"AWS Access Key ID\", None),\n (\"secret_access_key\", \"AWS Secret Access Key\", None),\n (\"region\", \"AWS Region\", None),\n ];\n\n fn chat_completions_builder(\n &self,\n client: &ReqwestClient,\n data: ChatCompletionsData,\n ) -> Result {\n let access_key_id = self.get_access_key_id()?;\n let secret_access_key = self.get_secret_access_key()?;\n let region = self.get_region()?;\n let session_token = self.get_session_token().ok();\n let host = format!(\"bedrock-runtime.{region}.amazonaws.com\");\n\n let model_name = &self.model.real_name();\n\n let uri = if data.stream {\n format!(\"/model/{model_name}/converse-stream\")\n } else {\n format!(\"/model/{model_name}/converse\")\n };\n\n let body = build_chat_completions_body(data, &self.model)?;\n\n let mut request_data = RequestData::new(\"\", body);\n self.patch_request_data(&mut request_data);\n let RequestData {\n url: _,\n headers,\n body,\n } = request_data;\n\n let builder = aws_fetch(\n client,\n &AwsCredentials {\n access_key_id,\n secret_access_key,\n region,\n session_token,\n },\n AwsRequest {\n method: Method::POST,\n host,\n service: \"bedrock\".into(),\n uri,\n querystring: \"\".into(),\n headers,\n body: body.to_string(),\n },\n )?;\n\n Ok(builder)\n }\n\n fn embeddings_builder(\n &self,\n client: &ReqwestClient,\n data: &EmbeddingsData,\n ) -> Result {\n let access_key_id = self.get_access_key_id()?;\n let secret_access_key = self.get_secret_access_key()?;\n let region = self.get_region()?;\n let session_token = self.get_session_token().ok();\n let host = format!(\"bedrock-runtime.{region}.amazonaws.com\");\n\n let uri = format!(\"/model/{}/invoke\", self.model.real_name());\n\n let input_type = match data.query {\n true => \"search_query\",\n false => \"search_document\",\n };\n\n let body = json!({\n \"texts\": data.texts,\n \"input_type\": input_type,\n });\n\n let mut request_data = RequestData::new(\"\", body);\n self.patch_request_data(&mut request_data);\n let RequestData {\n url: _,\n headers,\n body,\n } = request_data;\n\n let builder = aws_fetch(\n client,\n &AwsCredentials {\n access_key_id,\n secret_access_key,\n region,\n session_token,\n },\n AwsRequest {\n method: Method::POST,\n host,\n service: \"bedrock\".into(),\n uri,\n querystring: \"\".into(),\n headers,\n body: body.to_string(),\n },\n )?;\n\n Ok(builder)\n }\n}\n\n#[async_trait::async_trait]\nimpl Client for BedrockClient {\n client_common_fns!();\n\n async fn chat_completions_inner(\n &self,\n client: &ReqwestClient,\n data: ChatCompletionsData,\n ) -> Result {\n let builder = self.chat_completions_builder(client, data)?;\n chat_completions(builder).await\n }\n\n async fn chat_completions_streaming_inner(\n &self,\n client: &ReqwestClient,\n handler: &mut SseHandler,\n data: ChatCompletionsData,\n ) -> Result<()> {\n let builder = self.chat_completions_builder(client, data)?;\n chat_completions_streaming(builder, handler).await\n }\n\n async fn embeddings_inner(\n &self,\n client: &ReqwestClient,\n data: &EmbeddingsData,\n ) -> Result {\n let builder = self.embeddings_builder(client, data)?;\n embeddings(builder).await\n }\n}\n\nasync fn chat_completions(builder: RequestBuilder) -> Result {\n let res = builder.send().await?;\n let status = res.status();\n let data: Value = res.json().await?;\n\n if !status.is_success() {\n catch_error(&data, status.as_u16())?;\n }\n\n debug!(\"non-stream-data: {data}\");\n extract_chat_completions(&data)\n}\n\nasync fn chat_completions_streaming(\n builder: RequestBuilder,\n handler: &mut SseHandler,\n) -> Result<()> {\n let res = builder.send().await?;\n let status = res.status();\n if !status.is_success() {\n let data: Value = res.json().await?;\n catch_error(&data, status.as_u16())?;\n bail!(\"Invalid response data: {data}\");\n }\n\n let mut function_name = String::new();\n let mut function_arguments = String::new();\n let mut function_id = String::new();\n let mut reasoning_state = 0;\n\n let mut stream = res.bytes_stream();\n let mut buffer = BytesMut::new();\n let mut decoder = MessageFrameDecoder::new();\n while let Some(chunk) = stream.next().await {\n let chunk = chunk?;\n buffer.extend_from_slice(&chunk);\n while let DecodedFrame::Complete(message) = decoder.decode_frame(&mut buffer)? {\n let response_headers = parse_response_headers(&message)?;\n let message_type = response_headers.message_type.as_str();\n let smithy_type = response_headers.smithy_type.as_str();\n match (message_type, smithy_type) {\n (\"event\", _) => {\n let data: Value = serde_json::from_slice(message.payload())?;\n debug!(\"stream-data: {smithy_type} {data}\");\n match smithy_type {\n \"contentBlockStart\" => {\n if let Some(tool_use) = data[\"start\"][\"toolUse\"].as_object() {\n if let (Some(id), Some(name)) = (\n json_str_from_map(tool_use, \"toolUseId\"),\n json_str_from_map(tool_use, \"name\"),\n ) {\n if !function_name.is_empty() {\n if function_arguments.is_empty() {\n function_arguments = String::from(\"{}\");\n }\n let arguments: Value =\n function_arguments.parse().with_context(|| {\n format!(\"Tool call '{function_name}' have non-JSON arguments '{function_arguments}'\")\n })?;\n handler.tool_call(ToolCall::new(\n function_name.clone(),\n arguments,\n Some(function_id.clone()),\n ))?;\n }\n function_arguments.clear();\n function_name = name.into();\n function_id = id.into();\n }\n }\n }\n \"contentBlockDelta\" => {\n if let Some(text) = data[\"delta\"][\"text\"].as_str() {\n handler.text(text)?;\n } else if let Some(text) =\n data[\"delta\"][\"reasoningContent\"][\"text\"].as_str()\n {\n if reasoning_state == 0 {\n handler.text(\"\\n\")?;\n reasoning_state = 1;\n }\n handler.text(text)?;\n } else if let Some(input) = data[\"delta\"][\"toolUse\"][\"input\"].as_str() {\n function_arguments.push_str(input);\n }\n }\n \"contentBlockStop\" => {\n if reasoning_state == 1 {\n handler.text(\"\\n\\n\\n\")?;\n reasoning_state = 0;\n }\n if !function_name.is_empty() {\n if function_arguments.is_empty() {\n function_arguments = String::from(\"{}\");\n }\n let arguments: Value = function_arguments.parse().with_context(|| {\n format!(\"Tool call '{function_name}' have non-JSON arguments '{function_arguments}'\")\n })?;\n handler.tool_call(ToolCall::new(\n function_name.clone(),\n arguments,\n Some(function_id.clone()),\n ))?;\n }\n }\n _ => {}\n }\n }\n (\"exception\", _) => {\n let payload = base64_decode(message.payload())?;\n let data = String::from_utf8_lossy(&payload);\n\n bail!(\"Invalid response data: {data} (smithy_type: {smithy_type})\")\n }\n _ => {\n bail!(\"Unrecognized message, message_type: {message_type}, smithy_type: {smithy_type}\",);\n }\n }\n }\n }\n Ok(())\n}\n\nasync fn embeddings(builder: RequestBuilder) -> Result {\n let res = builder.send().await?;\n let status = res.status();\n let data: Value = res.json().await?;\n\n if !status.is_success() {\n catch_error(&data, status.as_u16())?;\n }\n\n let res_body: EmbeddingsResBody =\n serde_json::from_value(data).context(\"Invalid embeddings data\")?;\n Ok(res_body.embeddings)\n}\n\n#[derive(Deserialize)]\nstruct EmbeddingsResBody {\n embeddings: Vec>,\n}\n\nfn build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Result {\n let ChatCompletionsData {\n mut messages,\n temperature,\n top_p,\n functions,\n stream: _,\n } = data;\n\n let system_message = extract_system_message(&mut messages);\n\n let mut network_image_urls = vec![];\n\n let messages_len = messages.len();\n let messages: Vec = messages\n .into_iter()\n .enumerate()\n .flat_map(|(i, message)| {\n let Message { role, content } = message;\n match content {\n MessageContent::Text(text) if role.is_assistant() && i != messages_len - 1 => {\n vec![json!({ \"role\": role, \"content\": [ { \"text\": strip_think_tag(&text) } ] })]\n }\n MessageContent::Text(text) => vec![json!({\n \"role\": role,\n \"content\": [\n {\n \"text\": text,\n }\n ],\n })],\n MessageContent::Array(list) => {\n let content: Vec<_> = list\n .into_iter()\n .map(|item| match item {\n MessageContentPart::Text { text } => {\n json!({\"text\": text})\n }\n MessageContentPart::ImageUrl {\n image_url: ImageUrl { url },\n } => {\n if let Some((mime_type, data)) = url\n .strip_prefix(\"data:\")\n .and_then(|v| v.split_once(\";base64,\"))\n {\n json!({\n \"image\": {\n \"format\": mime_type.replace(\"image/\", \"\"),\n \"source\": {\n \"bytes\": data,\n }\n }\n })\n } else {\n network_image_urls.push(url.clone());\n json!({ \"url\": url })\n }\n }\n })\n .collect();\n vec![json!({\n \"role\": role,\n \"content\": content,\n })]\n }\n MessageContent::ToolCalls(MessageContentToolCalls {\n tool_results, text, ..\n }) => {\n let mut assistant_parts = vec![];\n let mut user_parts = vec![];\n if !text.is_empty() {\n assistant_parts.push(json!({\n \"text\": text,\n }))\n }\n for tool_result in tool_results {\n assistant_parts.push(json!({\n \"toolUse\": {\n \"toolUseId\": tool_result.call.id,\n \"name\": tool_result.call.name,\n \"input\": tool_result.call.arguments,\n }\n }));\n user_parts.push(json!({\n \"toolResult\": {\n \"toolUseId\": tool_result.call.id,\n \"content\": [\n {\n \"json\": tool_result.output,\n }\n ]\n }\n }));\n }\n vec![\n json!({\n \"role\": \"assistant\",\n \"content\": assistant_parts,\n }),\n json!({\n \"role\": \"user\",\n \"content\": user_parts,\n }),\n ]\n }\n }\n })\n .collect();\n\n if !network_image_urls.is_empty() {\n bail!(\n \"The model does not support network images: {:?}\",\n network_image_urls\n );\n }\n\n let mut body = json!({\n \"inferenceConfig\": {},\n \"messages\": messages,\n });\n if let Some(v) = system_message {\n body[\"system\"] = json!([\n {\n \"text\": v,\n }\n ])\n }\n\n if let Some(v) = model.max_tokens_param() {\n body[\"inferenceConfig\"][\"maxTokens\"] = v.into();\n }\n if let Some(v) = temperature {\n body[\"inferenceConfig\"][\"temperature\"] = v.into();\n }\n if let Some(v) = top_p {\n body[\"inferenceConfig\"][\"topP\"] = v.into();\n }\n if let Some(functions) = functions {\n let tools: Vec<_> = functions\n .iter()\n .map(|v| {\n json!({\n \"toolSpec\": {\n \"name\": v.name,\n \"description\": v.description,\n \"inputSchema\": {\n \"json\": v.parameters,\n },\n }\n })\n })\n .collect();\n body[\"toolConfig\"] = json!({\n \"tools\": tools,\n })\n }\n Ok(body)\n}\n\nfn extract_chat_completions(data: &Value) -> Result {\n let mut text = String::new();\n let mut reasoning = None;\n let mut tool_calls = vec![];\n if let Some(array) = data[\"output\"][\"message\"][\"content\"].as_array() {\n for item in array {\n if let Some(v) = item[\"text\"].as_str() {\n if !text.is_empty() {\n text.push_str(\"\\n\\n\");\n }\n text.push_str(v);\n } else if let Some(reasoning_text) =\n item[\"reasoningContent\"][\"reasoningText\"].as_object()\n {\n if let Some(text) = json_str_from_map(reasoning_text, \"text\") {\n reasoning = Some(text.to_string());\n }\n } else if let Some(tool_use) = item[\"toolUse\"].as_object() {\n if let (Some(id), Some(name), Some(input)) = (\n json_str_from_map(tool_use, \"toolUseId\"),\n json_str_from_map(tool_use, \"name\"),\n tool_use.get(\"input\"),\n ) {\n tool_calls.push(ToolCall::new(\n name.to_string(),\n input.clone(),\n Some(id.to_string()),\n ))\n }\n }\n }\n }\n\n if let Some(reasoning) = reasoning {\n text = format!(\"\\n{reasoning}\\n\\n\\n{text}\")\n }\n\n if text.is_empty() && tool_calls.is_empty() {\n bail!(\"Invalid response data: {data}\");\n }\n\n let output = ChatCompletionsOutput {\n text,\n tool_calls,\n id: None,\n input_tokens: data[\"usage\"][\"inputTokens\"].as_u64(),\n output_tokens: data[\"usage\"][\"outputTokens\"].as_u64(),\n };\n Ok(output)\n}\n\n#[derive(Debug)]\nstruct AwsCredentials {\n access_key_id: String,\n secret_access_key: String,\n region: String,\n session_token: Option,\n}\n\n#[derive(Debug)]\nstruct AwsRequest {\n method: Method,\n host: String,\n service: String,\n uri: String,\n querystring: String,\n headers: IndexMap,\n body: String,\n}\n\nfn aws_fetch(\n client: &ReqwestClient,\n credentials: &AwsCredentials,\n request: AwsRequest,\n) -> Result {\n let AwsRequest {\n method,\n host,\n service,\n uri,\n querystring,\n mut headers,\n body,\n } = request;\n let region = &credentials.region;\n\n let endpoint = format!(\"https://{host}{uri}\");\n\n let now: DateTime = Utc::now();\n let amz_date = now.format(\"%Y%m%dT%H%M%SZ\").to_string();\n let date_stamp = amz_date[0..8].to_string();\n headers.insert(\"host\".into(), host.clone());\n headers.insert(\"x-amz-date\".into(), amz_date.clone());\n if let Some(token) = credentials.session_token.clone() {\n headers.insert(\"x-amz-security-token\".into(), token);\n }\n\n let canonical_headers = headers\n .iter()\n .map(|(key, value)| format!(\"{key}:{value}\\n\"))\n .collect::>()\n .join(\"\");\n\n let signed_headers = headers\n .iter()\n .map(|(key, _)| key.as_str())\n .collect::>()\n .join(\";\");\n\n let payload_hash = sha256(&body);\n\n let canonical_request = format!(\n \"{}\\n{}\\n{}\\n{}\\n{}\\n{}\",\n method,\n encode_uri(&uri),\n querystring,\n canonical_headers,\n signed_headers,\n payload_hash\n );\n\n let algorithm = \"AWS4-HMAC-SHA256\";\n let credential_scope = format!(\"{date_stamp}/{region}/{service}/aws4_request\");\n let string_to_sign = format!(\n \"{}\\n{}\\n{}\\n{}\",\n algorithm,\n amz_date,\n credential_scope,\n sha256(&canonical_request)\n );\n\n let signing_key = gen_signing_key(\n &credentials.secret_access_key,\n &date_stamp,\n region,\n &service,\n );\n let signature = hmac_sha256(&signing_key, &string_to_sign);\n let signature = hex_encode(&signature);\n\n let authorization_header = format!(\n \"{} Credential={}/{}, SignedHeaders={}, Signature={}\",\n algorithm, credentials.access_key_id, credential_scope, signed_headers, signature\n );\n\n headers.insert(\"authorization\".into(), authorization_header);\n\n debug!(\"Request {endpoint} {body}\");\n\n let mut request_builder = client.request(method, endpoint).body(body);\n\n for (key, value) in &headers {\n request_builder = request_builder.header(key, value);\n }\n\n Ok(request_builder)\n}\n\nfn gen_signing_key(key: &str, date_stamp: &str, region: &str, service: &str) -> Vec {\n let k_date = hmac_sha256(format!(\"AWS4{key}\").as_bytes(), date_stamp);\n let k_region = hmac_sha256(&k_date, region);\n let k_service = hmac_sha256(&k_region, service);\n hmac_sha256(&k_service, \"aws4_request\")\n}\n" }, { "path": "src/client/claude.rs", "content": "use super::*;\n\nuse crate::utils::strip_think_tag;\n\nuse anyhow::{bail, Context, Result};\nuse reqwest::RequestBuilder;\nuse serde::Deserialize;\nuse serde_json::{json, Value};\n\nconst API_BASE: &str = \"https://api.anthropic.com/v1\";\n\n#[derive(Debug, Clone, Deserialize)]\npub struct ClaudeConfig {\n pub name: Option,\n pub api_key: Option,\n pub api_base: Option,\n #[serde(default)]\n pub models: Vec,\n pub patch: Option,\n pub extra: Option,\n}\n\nimpl ClaudeClient {\n config_get_fn!(api_key, get_api_key);\n config_get_fn!(api_base, get_api_base);\n\n pub const PROMPTS: [PromptAction<'static>; 1] = [(\"api_key\", \"API Key\", None)];\n}\n\nimpl_client_trait!(\n ClaudeClient,\n (\n prepare_chat_completions,\n claude_chat_completions,\n claude_chat_completions_streaming\n ),\n (noop_prepare_embeddings, noop_embeddings),\n (noop_prepare_rerank, noop_rerank),\n);\n\nfn prepare_chat_completions(\n self_: &ClaudeClient,\n data: ChatCompletionsData,\n) -> Result {\n let api_key = self_.get_api_key()?;\n let api_base = self_\n .get_api_base()\n .unwrap_or_else(|_| API_BASE.to_string());\n\n let url = format!(\"{}/messages\", api_base.trim_end_matches('/'));\n let body = claude_build_chat_completions_body(data, &self_.model)?;\n\n let mut request_data = RequestData::new(url, body);\n\n request_data.header(\"anthropic-version\", \"2023-06-01\");\n request_data.header(\"x-api-key\", api_key);\n\n Ok(request_data)\n}\n\npub async fn claude_chat_completions(\n builder: RequestBuilder,\n _model: &Model,\n) -> Result {\n let res = builder.send().await?;\n let status = res.status();\n let data: Value = res.json().await?;\n if !status.is_success() {\n catch_error(&data, status.as_u16())?;\n }\n debug!(\"non-stream-data: {data}\");\n claude_extract_chat_completions(&data)\n}\n\npub async fn claude_chat_completions_streaming(\n builder: RequestBuilder,\n handler: &mut SseHandler,\n _model: &Model,\n) -> Result<()> {\n let mut function_name = String::new();\n let mut function_arguments = String::new();\n let mut function_id = String::new();\n let mut reasoning_state = 0;\n let handle = |message: SseMmessage| -> Result {\n let data: Value = serde_json::from_str(&message.data)?;\n debug!(\"stream-data: {data}\");\n if let Some(typ) = data[\"type\"].as_str() {\n match typ {\n \"content_block_start\" => {\n if let (Some(\"tool_use\"), Some(name), Some(id)) = (\n data[\"content_block\"][\"type\"].as_str(),\n data[\"content_block\"][\"name\"].as_str(),\n data[\"content_block\"][\"id\"].as_str(),\n ) {\n if !function_name.is_empty() {\n let arguments: Value =\n function_arguments.parse().with_context(|| {\n format!(\"Tool call '{function_name}' have non-JSON arguments '{function_arguments}'\")\n })?;\n handler.tool_call(ToolCall::new(\n function_name.clone(),\n arguments,\n Some(function_id.clone()),\n ))?;\n }\n function_name = name.into();\n function_arguments.clear();\n function_id = id.into();\n }\n }\n \"content_block_delta\" => {\n if let Some(text) = data[\"delta\"][\"text\"].as_str() {\n handler.text(text)?;\n } else if let Some(text) = data[\"delta\"][\"thinking\"].as_str() {\n if reasoning_state == 0 {\n handler.text(\"\\n\")?;\n reasoning_state = 1;\n }\n handler.text(text)?;\n } else if let (true, Some(partial_json)) = (\n !function_name.is_empty(),\n data[\"delta\"][\"partial_json\"].as_str(),\n ) {\n function_arguments.push_str(partial_json);\n }\n }\n \"content_block_stop\" => {\n if reasoning_state == 1 {\n handler.text(\"\\n\\n\\n\")?;\n reasoning_state = 0;\n }\n if !function_name.is_empty() {\n let arguments: Value = if function_arguments.is_empty() {\n json!({})\n } else {\n function_arguments.parse().with_context(|| {\n format!(\"Tool call '{function_name}' have non-JSON arguments '{function_arguments}'\")\n })?\n };\n handler.tool_call(ToolCall::new(\n function_name.clone(),\n arguments,\n Some(function_id.clone()),\n ))?;\n }\n }\n _ => {}\n }\n }\n Ok(false)\n };\n\n sse_stream(builder, handle).await\n}\n\npub fn claude_build_chat_completions_body(\n data: ChatCompletionsData,\n model: &Model,\n) -> Result {\n let ChatCompletionsData {\n mut messages,\n temperature,\n top_p,\n functions,\n stream,\n } = data;\n\n let system_message = extract_system_message(&mut messages);\n\n let mut network_image_urls = vec![];\n\n let messages_len = messages.len();\n let messages: Vec = messages\n .into_iter()\n .enumerate()\n .flat_map(|(i, message)| {\n let Message { role, content } = message;\n match content {\n MessageContent::Text(text) if role.is_assistant() && i != messages_len - 1 => {\n vec![json!({ \"role\": role, \"content\": strip_think_tag(&text) })]\n }\n MessageContent::Text(text) => vec![json!({\n \"role\": role,\n \"content\": text,\n })],\n MessageContent::Array(list) => {\n let content: Vec<_> = list\n .into_iter()\n .map(|item| match item {\n MessageContentPart::Text { text } => {\n json!({\"type\": \"text\", \"text\": text})\n }\n MessageContentPart::ImageUrl {\n image_url: ImageUrl { url },\n } => {\n if let Some((mime_type, data)) = url\n .strip_prefix(\"data:\")\n .and_then(|v| v.split_once(\";base64,\"))\n {\n json!({\n \"type\": \"image\",\n \"source\": {\n \"type\": \"base64\",\n \"media_type\": mime_type,\n \"data\": data,\n }\n })\n } else {\n network_image_urls.push(url.clone());\n json!({ \"url\": url })\n }\n }\n })\n .collect();\n vec![json!({\n \"role\": role,\n \"content\": content,\n })]\n }\n MessageContent::ToolCalls(MessageContentToolCalls {\n tool_results, text, ..\n }) => {\n let mut assistant_parts = vec![];\n let mut user_parts = vec![];\n if !text.is_empty() {\n assistant_parts.push(json!({\n \"type\": \"text\",\n \"text\": text,\n }))\n }\n for tool_result in tool_results {\n assistant_parts.push(json!({\n \"type\": \"tool_use\",\n \"id\": tool_result.call.id,\n \"name\": tool_result.call.name,\n \"input\": tool_result.call.arguments,\n }));\n user_parts.push(json!({\n \"type\": \"tool_result\",\n \"tool_use_id\": tool_result.call.id,\n \"content\": tool_result.output.to_string(),\n }));\n }\n vec![\n json!({\n \"role\": \"assistant\",\n \"content\": assistant_parts,\n }),\n json!({\n \"role\": \"user\",\n \"content\": user_parts,\n }),\n ]\n }\n }\n })\n .collect();\n\n if !network_image_urls.is_empty() {\n bail!(\n \"The model does not support network images: {:?}\",\n network_image_urls\n );\n }\n\n let mut body = json!({\n \"model\": model.real_name(),\n \"messages\": messages,\n });\n if let Some(v) = system_message {\n body[\"system\"] = v.into();\n }\n if let Some(v) = model.max_tokens_param() {\n body[\"max_tokens\"] = v.into();\n }\n if let Some(v) = temperature {\n body[\"temperature\"] = v.into();\n }\n if let Some(v) = top_p {\n body[\"top_p\"] = v.into();\n }\n if stream {\n body[\"stream\"] = true.into();\n }\n if let Some(functions) = functions {\n body[\"tools\"] = functions\n .iter()\n .map(|v| {\n json!({\n \"name\": v.name,\n \"description\": v.description,\n \"input_schema\": v.parameters,\n })\n })\n .collect();\n }\n Ok(body)\n}\n\npub fn claude_extract_chat_completions(data: &Value) -> Result {\n let mut text = String::new();\n let mut reasoning = None;\n let mut tool_calls = vec![];\n if let Some(list) = data[\"content\"].as_array() {\n for item in list {\n match item[\"type\"].as_str() {\n Some(\"thinking\") => {\n if let Some(v) = item[\"thinking\"].as_str() {\n reasoning = Some(v.to_string());\n }\n }\n Some(\"text\") => {\n if let Some(v) = item[\"text\"].as_str() {\n if !text.is_empty() {\n text.push_str(\"\\n\\n\");\n }\n text.push_str(v);\n }\n }\n Some(\"tool_use\") => {\n if let (Some(name), Some(input), Some(id)) = (\n item[\"name\"].as_str(),\n item.get(\"input\"),\n item[\"id\"].as_str(),\n ) {\n tool_calls.push(ToolCall::new(\n name.to_string(),\n input.clone(),\n Some(id.to_string()),\n ));\n }\n }\n _ => {}\n }\n }\n }\n if let Some(reasoning) = reasoning {\n text = format!(\"\\n{reasoning}\\n\\n\\n{text}\")\n }\n\n if text.is_empty() && tool_calls.is_empty() {\n bail!(\"Invalid response data: {data}\");\n }\n\n let output = ChatCompletionsOutput {\n text: text.to_string(),\n tool_calls,\n id: data[\"id\"].as_str().map(|v| v.to_string()),\n input_tokens: data[\"usage\"][\"input_tokens\"].as_u64(),\n output_tokens: data[\"usage\"][\"output_tokens\"].as_u64(),\n };\n Ok(output)\n}\n" }, { "path": "src/client/cohere.rs", "content": "use super::openai::*;\nuse super::openai_compatible::*;\nuse super::*;\n\nuse anyhow::{bail, Context, Result};\nuse reqwest::RequestBuilder;\nuse serde::Deserialize;\nuse serde_json::{json, Value};\n\nconst API_BASE: &str = \"https://api.cohere.ai/v2\";\n\n#[derive(Debug, Clone, Deserialize, Default)]\npub struct CohereConfig {\n pub name: Option,\n pub api_key: Option,\n pub api_base: Option,\n #[serde(default)]\n pub models: Vec,\n pub patch: Option,\n pub extra: Option,\n}\n\nimpl CohereClient {\n config_get_fn!(api_key, get_api_key);\n config_get_fn!(api_base, get_api_base);\n\n pub const PROMPTS: [PromptAction<'static>; 1] = [(\"api_key\", \"API Key\", None)];\n}\n\nimpl_client_trait!(\n CohereClient,\n (\n prepare_chat_completions,\n chat_completions,\n chat_completions_streaming\n ),\n (prepare_embeddings, embeddings),\n (prepare_rerank, generic_rerank),\n);\n\nfn prepare_chat_completions(\n self_: &CohereClient,\n data: ChatCompletionsData,\n) -> Result {\n let api_key = self_.get_api_key()?;\n let api_base = self_\n .get_api_base()\n .unwrap_or_else(|_| API_BASE.to_string());\n\n let url = format!(\"{}/chat\", api_base.trim_end_matches('/'));\n let mut body = openai_build_chat_completions_body(data, &self_.model);\n if let Some(obj) = body.as_object_mut() {\n if let Some(top_p) = obj.remove(\"top_p\") {\n obj.insert(\"p\".to_string(), top_p);\n }\n }\n\n let mut request_data = RequestData::new(url, body);\n\n request_data.bearer_auth(api_key);\n\n Ok(request_data)\n}\n\nfn prepare_embeddings(self_: &CohereClient, data: &EmbeddingsData) -> Result {\n let api_key = self_.get_api_key()?;\n let api_base = self_\n .get_api_base()\n .unwrap_or_else(|_| API_BASE.to_string());\n\n let url = format!(\"{}/embed\", api_base.trim_end_matches('/'));\n\n let input_type = match data.query {\n true => \"search_query\",\n false => \"search_document\",\n };\n\n let body = json!({\n \"model\": self_.model.real_name(),\n \"texts\": data.texts,\n \"input_type\": input_type,\n \"embedding_types\": [\"float\"],\n });\n\n let mut request_data = RequestData::new(url, body);\n\n request_data.bearer_auth(api_key);\n\n Ok(request_data)\n}\n\nfn prepare_rerank(self_: &CohereClient, data: &RerankData) -> Result {\n let api_key = self_.get_api_key()?;\n let api_base = self_\n .get_api_base()\n .unwrap_or_else(|_| API_BASE.to_string());\n\n let url = format!(\"{}/rerank\", api_base.trim_end_matches('/'));\n let body = generic_build_rerank_body(data, &self_.model);\n\n let mut request_data = RequestData::new(url, body);\n\n request_data.bearer_auth(api_key);\n\n Ok(request_data)\n}\n\nasync fn chat_completions(\n builder: RequestBuilder,\n _model: &Model,\n) -> Result {\n let res = builder.send().await?;\n let status = res.status();\n let data: Value = res.json().await?;\n if !status.is_success() {\n catch_error(&data, status.as_u16())?;\n }\n\n debug!(\"non-stream-data: {data}\");\n extract_chat_completions(&data)\n}\n\nasync fn chat_completions_streaming(\n builder: RequestBuilder,\n handler: &mut SseHandler,\n _model: &Model,\n) -> Result<()> {\n let mut function_name = String::new();\n let mut function_arguments = String::new();\n let mut function_id = String::new();\n let handle = |message: SseMmessage| -> Result {\n if message.data == \"[DONE]\" {\n return Ok(true);\n }\n let data: Value = serde_json::from_str(&message.data)?;\n debug!(\"stream-data: {data}\");\n if let Some(typ) = data[\"type\"].as_str() {\n match typ {\n \"content-delta\" => {\n if let Some(text) = data[\"delta\"][\"message\"][\"content\"][\"text\"].as_str() {\n handler.text(text)?;\n }\n }\n \"tool-plan-delta\" => {\n if let Some(text) = data[\"delta\"][\"message\"][\"tool_plan\"].as_str() {\n handler.text(text)?;\n }\n }\n \"tool-call-start\" => {\n if let (Some(function), Some(id)) = (\n data[\"delta\"][\"message\"][\"tool_calls\"][\"function\"].as_object(),\n data[\"delta\"][\"message\"][\"tool_calls\"][\"id\"].as_str(),\n ) {\n if let Some(name) = function.get(\"name\").and_then(|v| v.as_str()) {\n function_name = name.to_string();\n }\n function_id = id.to_string();\n }\n }\n \"tool-call-delta\" => {\n if let Some(text) =\n data[\"delta\"][\"message\"][\"tool_calls\"][\"function\"][\"arguments\"].as_str()\n {\n function_arguments.push_str(text);\n }\n }\n \"tool-call-end\" => {\n if !function_name.is_empty() {\n let arguments: Value = function_arguments.parse().with_context(|| {\n format!(\"Tool call '{function_name}' have non-JSON arguments '{function_arguments}'\")\n })?;\n handler.tool_call(ToolCall::new(\n function_name.clone(),\n arguments,\n Some(function_id.clone()),\n ))?;\n }\n function_name.clear();\n function_arguments.clear();\n function_id.clear();\n }\n _ => {}\n }\n }\n Ok(false)\n };\n\n sse_stream(builder, handle).await\n}\n\nasync fn embeddings(builder: RequestBuilder, _model: &Model) -> Result {\n let res = builder.send().await?;\n let status = res.status();\n let data: Value = res.json().await?;\n if !status.is_success() {\n catch_error(&data, status.as_u16())?;\n }\n let res_body: EmbeddingsResBody =\n serde_json::from_value(data).context(\"Invalid embeddings data\")?;\n Ok(res_body.embeddings.float)\n}\n\n#[derive(Deserialize)]\nstruct EmbeddingsResBody {\n embeddings: EmbeddingsResBodyEmbeddings,\n}\n\n#[derive(Deserialize)]\nstruct EmbeddingsResBodyEmbeddings {\n float: Vec>,\n}\n\nfn extract_chat_completions(data: &Value) -> Result {\n let mut text = data[\"message\"][\"content\"][0][\"text\"]\n .as_str()\n .unwrap_or_default()\n .to_string();\n\n let mut tool_calls = vec![];\n if let Some(calls) = data[\"message\"][\"tool_calls\"].as_array() {\n if text.is_empty() {\n if let Some(tool_plain) = data[\"message\"][\"tool_plan\"].as_str() {\n text = tool_plain.to_string();\n }\n }\n for call in calls {\n if let (Some(name), Some(arguments), Some(id)) = (\n call[\"function\"][\"name\"].as_str(),\n call[\"function\"][\"arguments\"].as_str(),\n call[\"id\"].as_str(),\n ) {\n let arguments: Value = arguments.parse().with_context(|| {\n format!(\"Tool call '{name}' have non-JSON arguments '{arguments}'\")\n })?;\n tool_calls.push(ToolCall::new(\n name.to_string(),\n arguments,\n Some(id.to_string()),\n ));\n }\n }\n }\n\n if text.is_empty() && tool_calls.is_empty() {\n bail!(\"Invalid response data: {data}\");\n }\n let output = ChatCompletionsOutput {\n text,\n tool_calls,\n id: data[\"id\"].as_str().map(|v| v.to_string()),\n input_tokens: data[\"usage\"][\"billed_units\"][\"input_tokens\"].as_u64(),\n output_tokens: data[\"usage\"][\"billed_units\"][\"output_tokens\"].as_u64(),\n };\n Ok(output)\n}\n" }, { "path": "src/client/common.rs", "content": "use super::*;\n\nuse crate::{\n config::{Config, GlobalConfig, Input},\n function::{eval_tool_calls, FunctionDeclaration, ToolCall, ToolResult},\n render::render_stream,\n utils::*,\n};\n\nuse anyhow::{bail, Context, Result};\nuse fancy_regex::Regex;\nuse indexmap::IndexMap;\nuse inquire::{\n list_option::ListOption, required, validator::Validation, MultiSelect, Select, Text,\n};\nuse reqwest::{Client as ReqwestClient, RequestBuilder};\nuse serde::Deserialize;\nuse serde_json::{json, Value};\nuse std::sync::LazyLock;\nuse std::time::Duration;\nuse tokio::sync::mpsc::unbounded_channel;\n\nconst MODELS_YAML: &str = include_str!(\"../../models.yaml\");\n\npub static ALL_PROVIDER_MODELS: LazyLock> = LazyLock::new(|| {\n Config::loal_models_override()\n .ok()\n .unwrap_or_else(|| serde_yaml::from_str(MODELS_YAML).unwrap())\n});\n\nstatic EMBEDDING_MODEL_RE: LazyLock = LazyLock::new(|| {\n Regex::new(r\"((^|/)(bge-|e5-|uae-|gte-|text-)|embed|multilingual|minilm)\").unwrap()\n});\n\nstatic ESCAPE_SLASH_RE: LazyLock = LazyLock::new(|| Regex::new(r\"(? &GlobalConfig;\n\n fn extra_config(&self) -> Option<&ExtraConfig>;\n\n fn patch_config(&self) -> Option<&RequestPatch>;\n\n fn name(&self) -> &str;\n\n fn model(&self) -> &Model;\n\n fn model_mut(&mut self) -> &mut Model;\n\n fn build_client(&self) -> Result {\n let mut builder = ReqwestClient::builder();\n let extra = self.extra_config();\n let timeout = extra.and_then(|v| v.connect_timeout).unwrap_or(10);\n if let Some(proxy) = extra.and_then(|v| v.proxy.as_deref()) {\n builder = set_proxy(builder, proxy)?;\n }\n if let Some(user_agent) = self.global_config().read().user_agent.as_ref() {\n builder = builder.user_agent(user_agent);\n }\n let client = builder\n .connect_timeout(Duration::from_secs(timeout))\n .build()\n .with_context(|| \"Failed to build client\")?;\n Ok(client)\n }\n\n async fn chat_completions(&self, input: Input) -> Result {\n if self.global_config().read().dry_run {\n let content = input.echo_messages();\n return Ok(ChatCompletionsOutput::new(&content));\n }\n let client = self.build_client()?;\n let data = input.prepare_completion_data(self.model(), false)?;\n self.chat_completions_inner(&client, data)\n .await\n .with_context(|| \"Failed to call chat-completions api\")\n }\n\n async fn chat_completions_streaming(\n &self,\n input: &Input,\n handler: &mut SseHandler,\n ) -> Result<()> {\n let abort_signal = handler.abort();\n let input = input.clone();\n tokio::select! {\n ret = async {\n if self.global_config().read().dry_run {\n let content = input.echo_messages();\n handler.text(&content)?;\n return Ok(());\n }\n let client = self.build_client()?;\n let data = input.prepare_completion_data(self.model(), true)?;\n self.chat_completions_streaming_inner(&client, handler, data).await\n } => {\n handler.done();\n ret.with_context(|| \"Failed to call chat-completions api\")\n }\n _ = wait_abort_signal(&abort_signal) => {\n handler.done();\n Ok(())\n },\n }\n }\n\n async fn embeddings(&self, data: &EmbeddingsData) -> Result>> {\n let client = self.build_client()?;\n self.embeddings_inner(&client, data)\n .await\n .context(\"Failed to call embeddings api\")\n }\n\n async fn rerank(&self, data: &RerankData) -> Result {\n let client = self.build_client()?;\n self.rerank_inner(&client, data)\n .await\n .context(\"Failed to call rerank api\")\n }\n\n async fn chat_completions_inner(\n &self,\n client: &ReqwestClient,\n data: ChatCompletionsData,\n ) -> Result;\n\n async fn chat_completions_streaming_inner(\n &self,\n client: &ReqwestClient,\n handler: &mut SseHandler,\n data: ChatCompletionsData,\n ) -> Result<()>;\n\n async fn embeddings_inner(\n &self,\n _client: &ReqwestClient,\n _data: &EmbeddingsData,\n ) -> Result {\n bail!(\"The client doesn't support embeddings api\")\n }\n\n async fn rerank_inner(\n &self,\n _client: &ReqwestClient,\n _data: &RerankData,\n ) -> Result {\n bail!(\"The client doesn't support rerank api\")\n }\n\n fn request_builder(\n &self,\n client: &reqwest::Client,\n mut request_data: RequestData,\n ) -> RequestBuilder {\n self.patch_request_data(&mut request_data);\n request_data.into_builder(client)\n }\n\n fn patch_request_data(&self, request_data: &mut RequestData) {\n let model_type = self.model().model_type();\n if let Some(patch) = self.model().patch() {\n request_data.apply_patch(patch.clone());\n }\n\n let patch_map = std::env::var(get_env_name(&format!(\n \"patch_{}_{}\",\n self.model().client_name(),\n model_type.api_name(),\n )))\n .ok()\n .and_then(|v| serde_json::from_str(&v).ok())\n .or_else(|| {\n self.patch_config()\n .and_then(|v| model_type.extract_patch(v))\n .cloned()\n });\n let patch_map = match patch_map {\n Some(v) => v,\n _ => return,\n };\n for (key, patch) in patch_map {\n let key = ESCAPE_SLASH_RE.replace_all(&key, r\"\\/\");\n if let Ok(regex) = Regex::new(&format!(\"^({key})$\")) {\n if let Ok(true) = regex.is_match(self.model().name()) {\n request_data.apply_patch(patch);\n return;\n }\n }\n }\n }\n}\n\nimpl Default for ClientConfig {\n fn default() -> Self {\n Self::OpenAIConfig(OpenAIConfig::default())\n }\n}\n\n#[derive(Debug, Clone, Deserialize, Default)]\npub struct ExtraConfig {\n pub proxy: Option,\n pub connect_timeout: Option,\n}\n\n#[derive(Debug, Clone, Deserialize, Default)]\npub struct RequestPatch {\n pub chat_completions: Option,\n pub embeddings: Option,\n pub rerank: Option,\n}\n\npub type ApiPatch = IndexMap;\n\npub struct RequestData {\n pub url: String,\n pub headers: IndexMap,\n pub body: Value,\n}\n\nimpl RequestData {\n pub fn new(url: T, body: Value) -> Self\n where\n T: std::fmt::Display,\n {\n Self {\n url: url.to_string(),\n headers: Default::default(),\n body,\n }\n }\n\n pub fn bearer_auth(&mut self, auth: T)\n where\n T: std::fmt::Display,\n {\n self.headers\n .insert(\"authorization\".into(), format!(\"Bearer {auth}\"));\n }\n\n pub fn header(&mut self, key: K, value: V)\n where\n K: std::fmt::Display,\n V: std::fmt::Display,\n {\n self.headers.insert(key.to_string(), value.to_string());\n }\n\n pub fn into_builder(self, client: &ReqwestClient) -> RequestBuilder {\n let RequestData { url, headers, body } = self;\n debug!(\"Request {url} {body}\");\n\n let mut builder = client.post(url);\n for (key, value) in headers {\n builder = builder.header(key, value);\n }\n builder = builder.json(&body);\n builder\n }\n\n pub fn apply_patch(&mut self, patch: Value) {\n if let Some(patch_url) = patch[\"url\"].as_str() {\n self.url = patch_url.into();\n }\n if let Some(patch_body) = patch.get(\"body\") {\n json_patch::merge(&mut self.body, patch_body)\n }\n if let Some(patch_headers) = patch[\"headers\"].as_object() {\n for (key, value) in patch_headers {\n if let Some(value) = value.as_str() {\n self.header(key, value)\n } else if value.is_null() {\n self.headers.swap_remove(key);\n }\n }\n }\n }\n}\n\n#[derive(Debug)]\npub struct ChatCompletionsData {\n pub messages: Vec,\n pub temperature: Option,\n pub top_p: Option,\n pub functions: Option>,\n pub stream: bool,\n}\n\n#[derive(Debug, Clone, Default)]\npub struct ChatCompletionsOutput {\n pub text: String,\n pub tool_calls: Vec,\n pub id: Option,\n pub input_tokens: Option,\n pub output_tokens: Option,\n}\n\nimpl ChatCompletionsOutput {\n pub fn new(text: &str) -> Self {\n Self {\n text: text.to_string(),\n ..Default::default()\n }\n }\n}\n\n#[derive(Debug)]\npub struct EmbeddingsData {\n pub texts: Vec,\n pub query: bool,\n}\n\nimpl EmbeddingsData {\n pub fn new(texts: Vec, query: bool) -> Self {\n Self { texts, query }\n }\n}\n\npub type EmbeddingsOutput = Vec>;\n\n#[derive(Debug)]\npub struct RerankData {\n pub query: String,\n pub documents: Vec,\n pub top_n: usize,\n}\n\nimpl RerankData {\n pub fn new(query: String, documents: Vec, top_n: usize) -> Self {\n Self {\n query,\n documents,\n top_n,\n }\n }\n}\n\npub type RerankOutput = Vec;\n\n#[derive(Debug, Deserialize)]\npub struct RerankResult {\n pub index: usize,\n pub relevance_score: f64,\n}\n\npub type PromptAction<'a> = (&'a str, &'a str, Option<&'a str>);\n\npub async fn create_config(\n prompts: &[PromptAction<'static>],\n client: &str,\n) -> Result<(String, Value)> {\n let mut config = json!({\n \"type\": client,\n });\n for (key, desc, help_message) in prompts {\n let env_name = format!(\"{client}_{key}\").to_ascii_uppercase();\n let required = std::env::var(&env_name).is_err();\n let value = prompt_input_string(desc, required, *help_message)?;\n if !value.is_empty() {\n config[key] = value.into();\n }\n }\n let model = set_client_models_config(&mut config, client).await?;\n let clients = json!(vec![config]);\n Ok((model, clients))\n}\n\npub async fn create_openai_compatible_client_config(\n client: &str,\n) -> Result> {\n let api_base = super::OPENAI_COMPATIBLE_PROVIDERS\n .into_iter()\n .find(|(name, _)| client == *name)\n .map(|(_, api_base)| api_base)\n .unwrap_or(\"http(s)://{API_ADDR}/v1\");\n\n let name = if client == OpenAICompatibleClient::NAME {\n let value = prompt_input_string(\"Provider Name\", true, None)?;\n value.replace(' ', \"-\")\n } else {\n client.to_string()\n };\n\n let mut config = json!({\n \"type\": OpenAICompatibleClient::NAME,\n \"name\": &name,\n });\n\n let api_base = if api_base.contains('{') {\n prompt_input_string(\"API Base\", true, Some(&format!(\"e.g. {api_base}\")))?\n } else {\n api_base.to_string()\n };\n config[\"api_base\"] = api_base.into();\n\n let api_key = prompt_input_string(\"API Key\", false, None)?;\n if !api_key.is_empty() {\n config[\"api_key\"] = api_key.into();\n }\n\n let model = set_client_models_config(&mut config, &name).await?;\n let clients = json!(vec![config]);\n Ok(Some((model, clients)))\n}\n\npub async fn call_chat_completions(\n input: &Input,\n print: bool,\n extract_code: bool,\n client: &dyn Client,\n abort_signal: AbortSignal,\n) -> Result<(String, Vec)> {\n let ret = abortable_run_with_spinner(\n client.chat_completions(input.clone()),\n \"Generating\",\n abort_signal,\n )\n .await;\n\n match ret {\n Ok(ret) => {\n let ChatCompletionsOutput {\n mut text,\n tool_calls,\n ..\n } = ret;\n if !text.is_empty() {\n if extract_code {\n text = extract_code_block(&strip_think_tag(&text)).to_string();\n }\n if print {\n client.global_config().read().print_markdown(&text)?;\n }\n }\n Ok((text, eval_tool_calls(client.global_config(), tool_calls)?))\n }\n Err(err) => Err(err),\n }\n}\n\npub async fn call_chat_completions_streaming(\n input: &Input,\n client: &dyn Client,\n abort_signal: AbortSignal,\n) -> Result<(String, Vec)> {\n let (tx, rx) = unbounded_channel();\n let mut handler = SseHandler::new(tx, abort_signal.clone());\n\n let (send_ret, render_ret) = tokio::join!(\n client.chat_completions_streaming(input, &mut handler),\n render_stream(rx, client.global_config(), abort_signal.clone()),\n );\n\n if handler.abort().aborted() {\n bail!(\"Aborted.\");\n }\n\n render_ret?;\n\n let (text, tool_calls) = handler.take();\n match send_ret {\n Ok(_) => {\n if !text.is_empty() && !text.ends_with('\\n') {\n println!();\n }\n Ok((text, eval_tool_calls(client.global_config(), tool_calls)?))\n }\n Err(err) => {\n if !text.is_empty() {\n println!();\n }\n Err(err)\n }\n }\n}\n\npub fn noop_prepare_embeddings(_client: &T, _data: &EmbeddingsData) -> Result {\n bail!(\"The client doesn't support embeddings api\")\n}\n\npub async fn noop_embeddings(_builder: RequestBuilder, _model: &Model) -> Result {\n bail!(\"The client doesn't support embeddings api\")\n}\n\npub fn noop_prepare_rerank(_client: &T, _data: &RerankData) -> Result {\n bail!(\"The client doesn't support rerank api\")\n}\n\npub async fn noop_rerank(_builder: RequestBuilder, _model: &Model) -> Result {\n bail!(\"The client doesn't support rerank api\")\n}\n\npub fn catch_error(data: &Value, status: u16) -> Result<()> {\n if (200..300).contains(&status) {\n return Ok(());\n }\n debug!(\"Invalid response, status: {status}, data: {data}\");\n if let Some(error) = data[\"error\"].as_object() {\n if let (Some(typ), Some(message)) = (\n json_str_from_map(error, \"type\"),\n json_str_from_map(error, \"message\"),\n ) {\n bail!(\"{message} (type: {typ})\");\n } else if let (Some(typ), Some(message)) = (\n json_str_from_map(error, \"code\"),\n json_str_from_map(error, \"message\"),\n ) {\n bail!(\"{message} (code: {typ})\");\n }\n } else if let Some(error) = data[\"errors\"][0].as_object() {\n if let (Some(code), Some(message)) = (\n error.get(\"code\").and_then(|v| v.as_u64()),\n json_str_from_map(error, \"message\"),\n ) {\n bail!(\"{message} (status: {code})\")\n }\n } else if let Some(error) = data[0][\"error\"].as_object() {\n if let (Some(status), Some(message)) = (\n json_str_from_map(error, \"status\"),\n json_str_from_map(error, \"message\"),\n ) {\n bail!(\"{message} (status: {status})\")\n }\n } else if let (Some(detail), Some(status)) = (data[\"detail\"].as_str(), data[\"status\"].as_i64())\n {\n bail!(\"{detail} (status: {status})\");\n } else if let Some(error) = data[\"error\"].as_str() {\n bail!(\"{error}\");\n } else if let Some(message) = data[\"message\"].as_str() {\n bail!(\"{message}\");\n }\n bail!(\"Invalid response data: {data} (status: {status})\");\n}\n\npub fn json_str_from_map<'a>(\n map: &'a serde_json::Map,\n field_name: &str,\n) -> Option<&'a str> {\n map.get(field_name).and_then(|v| v.as_str())\n}\n\nasync fn set_client_models_config(client_config: &mut Value, client: &str) -> Result {\n if let Some(provider) = ALL_PROVIDER_MODELS.iter().find(|v| v.provider == client) {\n let models: Vec = provider\n .models\n .iter()\n .filter(|v| v.model_type == \"chat\")\n .map(|v| v.name.clone())\n .collect();\n let model_name = select_model(models)?;\n return Ok(format!(\"{client}:{model_name}\"));\n }\n let mut model_names = vec![];\n if let (Some(true), Some(api_base), api_key) = (\n client_config[\"type\"]\n .as_str()\n .map(|v| v == OpenAICompatibleClient::NAME),\n client_config[\"api_base\"].as_str(),\n client_config[\"api_key\"]\n .as_str()\n .map(|v| v.to_string())\n .or_else(|| {\n let env_name = format!(\"{client}_api_key\").to_ascii_uppercase();\n std::env::var(&env_name).ok()\n }),\n ) {\n match abortable_run_with_spinner(\n fetch_models(api_base, api_key.as_deref()),\n \"Fetching models\",\n create_abort_signal(),\n )\n .await\n {\n Ok(fetched_models) => {\n model_names = MultiSelect::new(\"LLMs to include (required):\", fetched_models)\n .with_validator(|list: &[ListOption<&String>]| {\n if list.is_empty() {\n Ok(Validation::Invalid(\n \"At least one item must be selected\".into(),\n ))\n } else {\n Ok(Validation::Valid)\n }\n })\n .prompt()?;\n }\n Err(err) => {\n eprintln!(\"✗ Fetch models failed: {err}\");\n }\n }\n }\n if model_names.is_empty() {\n model_names = prompt_input_string(\n \"LLMs to add\",\n true,\n Some(\"Separated by commas, e.g. llama3.3,qwen2.5\"),\n )?\n .split(',')\n .filter_map(|v| {\n let v = v.trim();\n if v.is_empty() {\n None\n } else {\n Some(v.to_string())\n }\n })\n .collect::>();\n }\n if model_names.is_empty() {\n bail!(\"No models\");\n }\n let models: Vec = model_names\n .iter()\n .map(|v| {\n let l = v.to_lowercase();\n if l.contains(\"rank\") {\n json!({\n \"name\": v,\n \"type\": \"reranker\",\n })\n } else if let Ok(true) = EMBEDDING_MODEL_RE.is_match(&l) {\n json!({\n \"name\": v,\n \"type\": \"embedding\",\n \"default_chunk_size\": 1000,\n \"max_batch_size\": 100\n })\n } else if v.contains(\"vision\") {\n json!({\n \"name\": v,\n \"supports_vision\": true\n })\n } else {\n json!({\n \"name\": v,\n })\n }\n })\n .collect();\n client_config[\"models\"] = models.into();\n let model_name = select_model(model_names)?;\n Ok(format!(\"{client}:{model_name}\"))\n}\n\nfn select_model(model_names: Vec) -> Result {\n if model_names.is_empty() {\n bail!(\"No models\");\n }\n let model = if model_names.len() == 1 {\n model_names[0].clone()\n } else {\n Select::new(\"Default Model (required):\", model_names).prompt()?\n };\n Ok(model)\n}\n\nfn prompt_input_string(\n desc: &str,\n required: bool,\n help_message: Option<&str>,\n) -> anyhow::Result {\n let desc = if required {\n format!(\"{desc} (required):\")\n } else {\n format!(\"{desc} (optional):\")\n };\n let mut text = Text::new(&desc);\n if required {\n text = text.with_validator(required!(\"This field is required\"))\n }\n if let Some(help_message) = help_message {\n text = text.with_help_message(help_message);\n }\n let text = text.prompt()?;\n Ok(text)\n}\n" }, { "path": "src/client/gemini.rs", "content": "use super::vertexai::*;\nuse super::*;\n\nuse anyhow::{Context, Result};\nuse reqwest::RequestBuilder;\nuse serde::Deserialize;\nuse serde_json::{json, Value};\n\nconst API_BASE: &str = \"https://generativelanguage.googleapis.com/v1beta\";\n\n#[derive(Debug, Clone, Deserialize, Default)]\npub struct GeminiConfig {\n pub name: Option,\n pub api_key: Option,\n pub api_base: Option,\n #[serde(default)]\n pub models: Vec,\n pub patch: Option,\n pub extra: Option,\n}\n\nimpl GeminiClient {\n config_get_fn!(api_key, get_api_key);\n config_get_fn!(api_base, get_api_base);\n\n pub const PROMPTS: [PromptAction<'static>; 1] = [(\"api_key\", \"API Key\", None)];\n}\n\nimpl_client_trait!(\n GeminiClient,\n (\n prepare_chat_completions,\n gemini_chat_completions,\n gemini_chat_completions_streaming\n ),\n (prepare_embeddings, embeddings),\n (noop_prepare_rerank, noop_rerank),\n);\n\nfn prepare_chat_completions(\n self_: &GeminiClient,\n data: ChatCompletionsData,\n) -> Result {\n let api_key = self_.get_api_key()?;\n let api_base = self_\n .get_api_base()\n .unwrap_or_else(|_| API_BASE.to_string());\n\n let func = match data.stream {\n true => \"streamGenerateContent\",\n false => \"generateContent\",\n };\n\n let url = format!(\n \"{}/models/{}:{}\",\n api_base.trim_end_matches('/'),\n self_.model.real_name(),\n func\n );\n\n let body = gemini_build_chat_completions_body(data, &self_.model)?;\n\n let mut request_data = RequestData::new(url, body);\n\n request_data.header(\"x-goog-api-key\", api_key);\n\n Ok(request_data)\n}\n\nfn prepare_embeddings(self_: &GeminiClient, data: &EmbeddingsData) -> Result {\n let api_key = self_.get_api_key()?;\n let api_base = self_\n .get_api_base()\n .unwrap_or_else(|_| API_BASE.to_string());\n\n let url = format!(\n \"{}/models/{}:batchEmbedContents?key={}\",\n api_base.trim_end_matches('/'),\n self_.model.real_name(),\n api_key\n );\n\n let model_id = format!(\"models/{}\", self_.model.real_name());\n\n let requests: Vec<_> = data\n .texts\n .iter()\n .map(|text| {\n json!({\n \"model\": model_id,\n \"content\": {\n \"parts\": [\n {\n \"text\": text\n }\n ]\n },\n })\n })\n .collect();\n\n let body = json!({\n \"requests\": requests,\n });\n\n let request_data = RequestData::new(url, body);\n\n Ok(request_data)\n}\n\nasync fn embeddings(builder: RequestBuilder, _model: &Model) -> Result {\n let res = builder.send().await?;\n let status = res.status();\n let data: Value = res.json().await?;\n if !status.is_success() {\n catch_error(&data, status.as_u16())?;\n }\n let res_body: EmbeddingsResBody =\n serde_json::from_value(data).context(\"Invalid embeddings data\")?;\n let output = res_body\n .embeddings\n .into_iter()\n .map(|embedding| embedding.values)\n .collect();\n Ok(output)\n}\n\n#[derive(Deserialize)]\nstruct EmbeddingsResBody {\n embeddings: Vec,\n}\n\n#[derive(Deserialize)]\nstruct EmbeddingsResBodyEmbedding {\n values: Vec,\n}\n" }, { "path": "src/client/macros.rs", "content": "#[macro_export]\nmacro_rules! register_client {\n (\n $(($module:ident, $name:literal, $config:ident, $client:ident),)+\n ) => {\n $(\n mod $module;\n )+\n $(\n use self::$module::$config;\n )+\n\n #[derive(Debug, Clone, serde::Deserialize)]\n #[serde(tag = \"type\")]\n pub enum ClientConfig {\n $(\n #[serde(rename = $name)]\n $config($config),\n )+\n #[serde(other)]\n Unknown,\n }\n\n $(\n #[derive(Debug)]\n pub struct $client {\n global_config: $crate::config::GlobalConfig,\n config: $config,\n model: $crate::client::Model,\n }\n\n impl $client {\n pub const NAME: &'static str = $name;\n\n pub fn init(global_config: &$crate::config::GlobalConfig, model: &$crate::client::Model) -> Option> {\n let config = global_config.read().clients.iter().find_map(|client_config| {\n if let ClientConfig::$config(c) = client_config {\n if Self::name(c) == model.client_name() {\n return Some(c.clone())\n }\n }\n None\n })?;\n\n Some(Box::new(Self {\n global_config: global_config.clone(),\n config,\n model: model.clone(),\n }))\n }\n\n pub fn list_models(local_config: &$config) -> Vec {\n let client_name = Self::name(local_config);\n if local_config.models.is_empty() {\n if let Some(v) = $crate::client::ALL_PROVIDER_MODELS.iter().find(|v| {\n v.provider == $name ||\n ($name == OpenAICompatibleClient::NAME\n && local_config.name.as_ref().map(|name| name.starts_with(&v.provider)).unwrap_or_default())\n }) {\n return Model::from_config(client_name, &v.models);\n }\n vec![]\n } else {\n Model::from_config(client_name, &local_config.models)\n }\n }\n\n pub fn name(local_config: &$config) -> &str {\n local_config.name.as_deref().unwrap_or(Self::NAME)\n }\n }\n\n )+\n\n pub fn init_client(config: &$crate::config::GlobalConfig, model: Option<$crate::client::Model>) -> anyhow::Result> {\n let model = model.unwrap_or_else(|| config.read().model.clone());\n None\n $(.or_else(|| $client::init(config, &model)))+\n .ok_or_else(|| {\n anyhow::anyhow!(\"Invalid model '{}'\", model.id())\n })\n }\n\n pub fn list_client_types() -> Vec<&'static str> {\n let mut client_types: Vec<_> = vec![$($client::NAME,)+];\n client_types.extend($crate::client::OPENAI_COMPATIBLE_PROVIDERS.iter().map(|(name, _)| *name));\n client_types\n }\n\n pub async fn create_client_config(client: &str) -> anyhow::Result<(String, serde_json::Value)> {\n $(\n if client == $client::NAME && client != $crate::client::OpenAICompatibleClient::NAME {\n return create_config(&$client::PROMPTS, $client::NAME).await\n }\n )+\n if let Some(ret) = create_openai_compatible_client_config(client).await? {\n return Ok(ret);\n }\n anyhow::bail!(\"Unknown client '{}'\", client)\n }\n\n static ALL_CLIENT_NAMES: std::sync::OnceLock> = std::sync::OnceLock::new();\n\n pub fn list_client_names(config: &$crate::config::Config) -> Vec<&'static String> {\n let names = ALL_CLIENT_NAMES.get_or_init(|| {\n config\n .clients\n .iter()\n .flat_map(|v| match v {\n $(ClientConfig::$config(c) => vec![$client::name(c).to_string()],)+\n ClientConfig::Unknown => vec![],\n })\n .collect()\n });\n names.iter().collect()\n }\n\n static ALL_MODELS: std::sync::OnceLock> = std::sync::OnceLock::new();\n\n pub fn list_all_models(config: &$crate::config::Config) -> Vec<&'static $crate::client::Model> {\n let models = ALL_MODELS.get_or_init(|| {\n config\n .clients\n .iter()\n .flat_map(|v| match v {\n $(ClientConfig::$config(c) => $client::list_models(c),)+\n ClientConfig::Unknown => vec![],\n })\n .collect()\n });\n models.iter().collect()\n }\n\n pub fn list_models(config: &$crate::config::Config, model_type: $crate::client::ModelType) -> Vec<&'static $crate::client::Model> {\n list_all_models(config).into_iter().filter(|v| v.model_type() == model_type).collect()\n }\n };\n}\n\n#[macro_export]\nmacro_rules! client_common_fns {\n () => {\n fn global_config(&self) -> &$crate::config::GlobalConfig {\n &self.global_config\n }\n\n fn extra_config(&self) -> Option<&$crate::client::ExtraConfig> {\n self.config.extra.as_ref()\n }\n\n fn patch_config(&self) -> Option<&$crate::client::RequestPatch> {\n self.config.patch.as_ref()\n }\n\n fn name(&self) -> &str {\n Self::name(&self.config)\n }\n\n fn model(&self) -> &Model {\n &self.model\n }\n\n fn model_mut(&mut self) -> &mut Model {\n &mut self.model\n }\n };\n}\n\n#[macro_export]\nmacro_rules! impl_client_trait {\n (\n $client:ident,\n ($prepare_chat_completions:path, $chat_completions:path, $chat_completions_streaming:path),\n ($prepare_embeddings:path, $embeddings:path),\n ($prepare_rerank:path, $rerank:path),\n ) => {\n #[async_trait::async_trait]\n impl $crate::client::Client for $crate::client::$client {\n client_common_fns!();\n\n async fn chat_completions_inner(\n &self,\n client: &reqwest::Client,\n data: $crate::client::ChatCompletionsData,\n ) -> anyhow::Result<$crate::client::ChatCompletionsOutput> {\n let request_data = $prepare_chat_completions(self, data)?;\n let builder = self.request_builder(client, request_data);\n $chat_completions(builder, self.model()).await\n }\n\n async fn chat_completions_streaming_inner(\n &self,\n client: &reqwest::Client,\n handler: &mut $crate::client::SseHandler,\n data: $crate::client::ChatCompletionsData,\n ) -> Result<()> {\n let request_data = $prepare_chat_completions(self, data)?;\n let builder = self.request_builder(client, request_data);\n $chat_completions_streaming(builder, handler, self.model()).await\n }\n\n async fn embeddings_inner(\n &self,\n client: &reqwest::Client,\n data: &$crate::client::EmbeddingsData,\n ) -> Result<$crate::client::EmbeddingsOutput> {\n let request_data = $prepare_embeddings(self, data)?;\n let builder = self.request_builder(client, request_data);\n $embeddings(builder, self.model()).await\n }\n\n async fn rerank_inner(\n &self,\n client: &reqwest::Client,\n data: &$crate::client::RerankData,\n ) -> Result<$crate::client::RerankOutput> {\n let request_data = $prepare_rerank(self, data)?;\n let builder = self.request_builder(client, request_data);\n $rerank(builder, self.model()).await\n }\n }\n };\n}\n\n#[macro_export]\nmacro_rules! config_get_fn {\n ($field_name:ident, $fn_name:ident) => {\n fn $fn_name(&self) -> anyhow::Result {\n let env_prefix = Self::name(&self.config);\n let env_name =\n format!(\"{}_{}\", env_prefix, stringify!($field_name)).to_ascii_uppercase();\n std::env::var(&env_name)\n .ok()\n .or_else(|| self.config.$field_name.clone())\n .ok_or_else(|| anyhow::anyhow!(\"Miss '{}'\", stringify!($field_name)))\n }\n };\n}\n\n#[macro_export]\nmacro_rules! unsupported_model {\n ($name:expr) => {\n anyhow::bail!(\"Unsupported model '{}'\", $name)\n };\n}\n" }, { "path": "src/client/message.rs", "content": "use super::Model;\n\nuse crate::{function::ToolResult, multiline_text, utils::dimmed_text};\n\nuse serde::{Deserialize, Serialize};\n\n#[derive(Debug, Clone, Deserialize, Serialize)]\npub struct Message {\n pub role: MessageRole,\n pub content: MessageContent,\n}\n\nimpl Default for Message {\n fn default() -> Self {\n Self {\n role: MessageRole::User,\n content: MessageContent::Text(String::new()),\n }\n }\n}\n\nimpl Message {\n pub fn new(role: MessageRole, content: MessageContent) -> Self {\n Self { role, content }\n }\n\n pub fn merge_system(&mut self, system: MessageContent) {\n match (&mut self.content, system) {\n (MessageContent::Text(text), MessageContent::Text(system_text)) => {\n self.content = MessageContent::Array(vec![\n MessageContentPart::Text { text: system_text },\n MessageContentPart::Text {\n text: text.to_string(),\n },\n ])\n }\n (MessageContent::Array(list), MessageContent::Text(system_text)) => {\n list.insert(0, MessageContentPart::Text { text: system_text })\n }\n (MessageContent::Text(text), MessageContent::Array(mut system_list)) => {\n system_list.push(MessageContentPart::Text {\n text: text.to_string(),\n });\n self.content = MessageContent::Array(system_list);\n }\n (MessageContent::Array(list), MessageContent::Array(mut system_list)) => {\n system_list.append(list);\n self.content = MessageContent::Array(system_list);\n }\n _ => {}\n }\n }\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]\n#[serde(rename_all = \"snake_case\")]\npub enum MessageRole {\n System,\n Assistant,\n User,\n Tool,\n}\n\n#[allow(dead_code)]\nimpl MessageRole {\n pub fn is_system(&self) -> bool {\n matches!(self, MessageRole::System)\n }\n\n pub fn is_user(&self) -> bool {\n matches!(self, MessageRole::User)\n }\n\n pub fn is_assistant(&self) -> bool {\n matches!(self, MessageRole::Assistant)\n }\n}\n\n#[derive(Debug, Clone, Deserialize, Serialize)]\n#[serde(untagged)]\npub enum MessageContent {\n Text(String),\n Array(Vec),\n // Note: This type is primarily for convenience and does not exist in OpenAI's API.\n ToolCalls(MessageContentToolCalls),\n}\n\nimpl MessageContent {\n pub fn render_input(\n &self,\n resolve_url_fn: impl Fn(&str) -> String,\n agent_info: &Option<(String, Vec)>,\n ) -> String {\n match self {\n MessageContent::Text(text) => multiline_text(text),\n MessageContent::Array(list) => {\n let (mut concated_text, mut files) = (String::new(), vec![]);\n for item in list {\n match item {\n MessageContentPart::Text { text } => {\n concated_text = format!(\"{concated_text} {text}\")\n }\n MessageContentPart::ImageUrl { image_url } => {\n files.push(resolve_url_fn(&image_url.url))\n }\n }\n }\n if !concated_text.is_empty() {\n concated_text = format!(\" -- {}\", multiline_text(&concated_text))\n }\n format!(\".file {}{}\", files.join(\" \"), concated_text)\n }\n MessageContent::ToolCalls(MessageContentToolCalls {\n tool_results, text, ..\n }) => {\n let mut lines = vec![];\n if !text.is_empty() {\n lines.push(text.clone())\n }\n for tool_result in tool_results {\n let mut parts = vec![\"Call\".to_string()];\n if let Some((agent_name, functions)) = agent_info {\n if functions.contains(&tool_result.call.name) {\n parts.push(agent_name.clone())\n }\n }\n parts.push(tool_result.call.name.clone());\n parts.push(tool_result.call.arguments.to_string());\n lines.push(dimmed_text(&parts.join(\" \")));\n }\n lines.join(\"\\n\")\n }\n }\n }\n\n pub fn merge_prompt(&mut self, replace_fn: impl Fn(&str) -> String) {\n match self {\n MessageContent::Text(text) => *text = replace_fn(text),\n MessageContent::Array(list) => {\n if list.is_empty() {\n list.push(MessageContentPart::Text {\n text: replace_fn(\"\"),\n })\n } else if let Some(MessageContentPart::Text { text }) = list.get_mut(0) {\n *text = replace_fn(text)\n }\n }\n MessageContent::ToolCalls(_) => {}\n }\n }\n\n pub fn to_text(&self) -> String {\n match self {\n MessageContent::Text(text) => text.to_string(),\n MessageContent::Array(list) => {\n let mut parts = vec![];\n for item in list {\n if let MessageContentPart::Text { text } = item {\n parts.push(text.clone())\n }\n }\n parts.join(\"\\n\\n\")\n }\n MessageContent::ToolCalls(_) => String::new(),\n }\n }\n}\n\n#[derive(Debug, Clone, Deserialize, Serialize)]\n#[serde(tag = \"type\", rename_all = \"snake_case\")]\npub enum MessageContentPart {\n Text { text: String },\n ImageUrl { image_url: ImageUrl },\n}\n\n#[derive(Debug, Clone, Deserialize, Serialize)]\npub struct ImageUrl {\n pub url: String,\n}\n\n#[derive(Debug, Clone, Deserialize, Serialize)]\npub struct MessageContentToolCalls {\n pub tool_results: Vec,\n pub text: String,\n pub sequence: bool,\n}\n\nimpl MessageContentToolCalls {\n pub fn new(tool_results: Vec, text: String) -> Self {\n Self {\n tool_results,\n text,\n sequence: false,\n }\n }\n\n pub fn merge(&mut self, tool_results: Vec, _text: String) {\n self.tool_results.extend(tool_results);\n self.text.clear();\n self.sequence = true;\n }\n}\n\npub fn patch_messages(messages: &mut Vec, model: &Model) {\n if messages.is_empty() {\n return;\n }\n if let Some(prefix) = model.system_prompt_prefix() {\n if messages[0].role.is_system() {\n messages[0].merge_system(MessageContent::Text(prefix.to_string()));\n } else {\n messages.insert(\n 0,\n Message {\n role: MessageRole::System,\n content: MessageContent::Text(prefix.to_string()),\n },\n );\n }\n }\n if model.no_system_message() && messages[0].role.is_system() {\n let system_message = messages.remove(0);\n if let (Some(message), system) = (messages.get_mut(0), system_message.content) {\n message.merge_system(system);\n }\n }\n}\n\npub fn extract_system_message(messages: &mut Vec) -> Option {\n if messages[0].role.is_system() {\n let system_message = messages.remove(0);\n return Some(system_message.content.to_text());\n }\n None\n}\n" }, { "path": "src/client/mod.rs", "content": "mod access_token;\nmod common;\nmod message;\n#[macro_use]\nmod macros;\nmod model;\nmod stream;\n\npub use crate::function::ToolCall;\npub use common::*;\npub use message::*;\npub use model::*;\npub use stream::*;\n\nregister_client!(\n (openai, \"openai\", OpenAIConfig, OpenAIClient),\n (\n openai_compatible,\n \"openai-compatible\",\n OpenAICompatibleConfig,\n OpenAICompatibleClient\n ),\n (gemini, \"gemini\", GeminiConfig, GeminiClient),\n (claude, \"claude\", ClaudeConfig, ClaudeClient),\n (cohere, \"cohere\", CohereConfig, CohereClient),\n (\n azure_openai,\n \"azure-openai\",\n AzureOpenAIConfig,\n AzureOpenAIClient\n ),\n (vertexai, \"vertexai\", VertexAIConfig, VertexAIClient),\n (bedrock, \"bedrock\", BedrockConfig, BedrockClient),\n);\n\npub const OPENAI_COMPATIBLE_PROVIDERS: [(&str, &str); 18] = [\n (\"ai21\", \"https://api.ai21.com/studio/v1\"),\n (\n \"cloudflare\",\n \"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/v1\",\n ),\n (\"deepinfra\", \"https://api.deepinfra.com/v1/openai\"),\n (\"deepseek\", \"https://api.deepseek.com\"),\n (\"ernie\", \"https://qianfan.baidubce.com/v2\"),\n (\"github\", \"https://models.inference.ai.azure.com\"),\n (\"groq\", \"https://api.groq.com/openai/v1\"),\n (\"hunyuan\", \"https://api.hunyuan.cloud.tencent.com/v1\"),\n (\"minimax\", \"https://api.minimax.chat/v1\"),\n (\"mistral\", \"https://api.mistral.ai/v1\"),\n (\"moonshot\", \"https://api.moonshot.cn/v1\"),\n (\"openrouter\", \"https://openrouter.ai/api/v1\"),\n (\"perplexity\", \"https://api.perplexity.ai\"),\n (\n \"qianwen\",\n \"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n ),\n (\"xai\", \"https://api.x.ai/v1\"),\n (\"zhipuai\", \"https://open.bigmodel.cn/api/paas/v4\"),\n // RAG-dedicated\n (\"jina\", \"https://api.jina.ai/v1\"),\n (\"voyageai\", \"https://api.voyageai.com/v1\"),\n];\n" }, { "path": "src/client/model.rs", "content": "use super::{\n list_all_models, list_client_names,\n message::{Message, MessageContent, MessageContentPart},\n ApiPatch, MessageContentToolCalls, RequestPatch,\n};\n\nuse crate::config::Config;\nuse crate::utils::{estimate_token_length, strip_think_tag};\n\nuse anyhow::{bail, Result};\nuse serde::{Deserialize, Serialize};\nuse serde_json::Value;\nuse std::fmt::Display;\n\nconst PER_MESSAGES_TOKENS: usize = 5;\nconst BASIS_TOKENS: usize = 2;\n\n#[derive(Debug, Clone)]\npub struct Model {\n client_name: String,\n data: ModelData,\n}\n\nimpl Default for Model {\n fn default() -> Self {\n Model::new(\"\", \"\")\n }\n}\n\nimpl Model {\n pub fn new(client_name: &str, name: &str) -> Self {\n Self {\n client_name: client_name.into(),\n data: ModelData::new(name),\n }\n }\n\n pub fn from_config(client_name: &str, models: &[ModelData]) -> Vec {\n models\n .iter()\n .map(|v| Model {\n client_name: client_name.to_string(),\n data: v.clone(),\n })\n .collect()\n }\n\n pub fn retrieve_model(config: &Config, model_id: &str, model_type: ModelType) -> Result {\n let models = list_all_models(config);\n let (client_name, model_name) = match model_id.split_once(':') {\n Some((client_name, model_name)) => {\n if model_name.is_empty() {\n (client_name, None)\n } else {\n (client_name, Some(model_name))\n }\n }\n None => (model_id, None),\n };\n match model_name {\n Some(model_name) => {\n if let Some(model) = models.iter().find(|v| v.id() == model_id) {\n if model.model_type() == model_type {\n return Ok((*model).clone());\n } else {\n bail!(\"Model '{model_id}' is not a {model_type} model\")\n }\n }\n if list_client_names(config)\n .into_iter()\n .any(|v| *v == client_name)\n && model_type.can_create_from_name()\n {\n let mut new_model = Self::new(client_name, model_name);\n new_model.data.model_type = model_type.to_string();\n return Ok(new_model);\n }\n }\n None => {\n if let Some(found) = models\n .iter()\n .find(|v| v.client_name == client_name && v.model_type() == model_type)\n {\n return Ok((*found).clone());\n }\n }\n };\n bail!(\"Unknown {model_type} model '{model_id}'\")\n }\n\n pub fn id(&self) -> String {\n if self.data.name.is_empty() {\n self.client_name.to_string()\n } else {\n format!(\"{}:{}\", self.client_name, self.data.name)\n }\n }\n\n pub fn client_name(&self) -> &str {\n &self.client_name\n }\n\n pub fn name(&self) -> &str {\n &self.data.name\n }\n\n pub fn real_name(&self) -> &str {\n self.data.real_name.as_deref().unwrap_or(&self.data.name)\n }\n\n pub fn model_type(&self) -> ModelType {\n if self.data.model_type.starts_with(\"embed\") {\n ModelType::Embedding\n } else if self.data.model_type.starts_with(\"rerank\") {\n ModelType::Reranker\n } else {\n ModelType::Chat\n }\n }\n\n pub fn data(&self) -> &ModelData {\n &self.data\n }\n\n pub fn data_mut(&mut self) -> &mut ModelData {\n &mut self.data\n }\n\n pub fn description(&self) -> String {\n match self.model_type() {\n ModelType::Chat => {\n let ModelData {\n max_input_tokens,\n max_output_tokens,\n input_price,\n output_price,\n supports_vision,\n supports_function_calling,\n ..\n } = &self.data;\n let max_input_tokens = stringify_option_value(max_input_tokens);\n let max_output_tokens = stringify_option_value(max_output_tokens);\n let input_price = stringify_option_value(input_price);\n let output_price = stringify_option_value(output_price);\n let mut capabilities = vec![];\n if *supports_vision {\n capabilities.push('👁');\n };\n if *supports_function_calling {\n capabilities.push('⚒');\n };\n let capabilities: String = capabilities\n .into_iter()\n .map(|v| format!(\"{v} \"))\n .collect::>()\n .join(\"\");\n format!(\n \"{max_input_tokens:>8} / {max_output_tokens:>8} | {input_price:>6} / {output_price:>6} {capabilities:>6}\"\n )\n }\n ModelType::Embedding => {\n let ModelData {\n input_price,\n max_tokens_per_chunk,\n max_batch_size,\n ..\n } = &self.data;\n let max_tokens = stringify_option_value(max_tokens_per_chunk);\n let max_batch = stringify_option_value(max_batch_size);\n let price = stringify_option_value(input_price);\n format!(\"max-tokens:{max_tokens};max-batch:{max_batch};price:{price}\")\n }\n ModelType::Reranker => String::new(),\n }\n }\n\n pub fn patch(&self) -> Option<&Value> {\n self.data.patch.as_ref()\n }\n\n pub fn max_input_tokens(&self) -> Option {\n self.data.max_input_tokens\n }\n\n pub fn max_output_tokens(&self) -> Option {\n self.data.max_output_tokens\n }\n\n pub fn no_stream(&self) -> bool {\n self.data.no_stream\n }\n\n pub fn no_system_message(&self) -> bool {\n self.data.no_system_message\n }\n\n pub fn system_prompt_prefix(&self) -> Option<&str> {\n self.data.system_prompt_prefix.as_deref()\n }\n\n pub fn max_tokens_per_chunk(&self) -> Option {\n self.data.max_tokens_per_chunk\n }\n\n pub fn default_chunk_size(&self) -> usize {\n self.data.default_chunk_size.unwrap_or(1000)\n }\n\n pub fn max_batch_size(&self) -> Option {\n self.data.max_batch_size\n }\n\n pub fn max_tokens_param(&self) -> Option {\n if self.data.require_max_tokens {\n self.data.max_output_tokens\n } else {\n None\n }\n }\n\n pub fn set_max_tokens(\n &mut self,\n max_output_tokens: Option,\n require_max_tokens: bool,\n ) -> &mut Self {\n match max_output_tokens {\n None | Some(0) => self.data.max_output_tokens = None,\n _ => self.data.max_output_tokens = max_output_tokens,\n }\n self.data.require_max_tokens = require_max_tokens;\n self\n }\n\n pub fn messages_tokens(&self, messages: &[Message]) -> usize {\n let messages_len = messages.len();\n messages\n .iter()\n .enumerate()\n .map(|(i, v)| match &v.content {\n MessageContent::Text(text) => {\n if v.role.is_assistant() && i != messages_len - 1 {\n estimate_token_length(&strip_think_tag(text))\n } else {\n estimate_token_length(text)\n }\n }\n MessageContent::Array(list) => list\n .iter()\n .map(|v| match v {\n MessageContentPart::Text { text } => estimate_token_length(text),\n MessageContentPart::ImageUrl { .. } => 0,\n })\n .sum(),\n MessageContent::ToolCalls(MessageContentToolCalls {\n tool_results, text, ..\n }) => {\n estimate_token_length(text)\n + tool_results\n .iter()\n .map(|v| {\n serde_json::to_string(v)\n .map(|v| estimate_token_length(&v))\n .unwrap_or_default()\n })\n .sum::()\n }\n })\n .sum()\n }\n\n pub fn total_tokens(&self, messages: &[Message]) -> usize {\n if messages.is_empty() {\n return 0;\n }\n let num_messages = messages.len();\n let message_tokens = self.messages_tokens(messages);\n if messages[num_messages - 1].role.is_user() {\n num_messages * PER_MESSAGES_TOKENS + message_tokens\n } else {\n (num_messages - 1) * PER_MESSAGES_TOKENS + message_tokens\n }\n }\n\n pub fn guard_max_input_tokens(&self, messages: &[Message]) -> Result<()> {\n let total_tokens = self.total_tokens(messages) + BASIS_TOKENS;\n if let Some(max_input_tokens) = self.data.max_input_tokens {\n if total_tokens >= max_input_tokens {\n bail!(\"Exceed max_input_tokens limit\")\n }\n }\n Ok(())\n }\n}\n\n#[derive(Debug, Clone, Default, Serialize, Deserialize)]\npub struct ModelData {\n pub name: String,\n #[serde(default = \"default_model_type\", rename = \"type\")]\n pub model_type: String,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub real_name: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub max_input_tokens: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub input_price: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub output_price: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub patch: Option,\n\n // chat-only properties\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub max_output_tokens: Option,\n #[serde(default, skip_serializing_if = \"std::ops::Not::not\")]\n pub require_max_tokens: bool,\n #[serde(default, skip_serializing_if = \"std::ops::Not::not\")]\n pub supports_vision: bool,\n #[serde(default, skip_serializing_if = \"std::ops::Not::not\")]\n pub supports_function_calling: bool,\n #[serde(default, skip_serializing_if = \"std::ops::Not::not\")]\n no_stream: bool,\n #[serde(default, skip_serializing_if = \"std::ops::Not::not\")]\n no_system_message: bool,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n system_prompt_prefix: Option,\n\n // embedding-only properties\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub max_tokens_per_chunk: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub default_chunk_size: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub max_batch_size: Option,\n}\n\nimpl ModelData {\n pub fn new(name: &str) -> Self {\n Self {\n name: name.to_string(),\n model_type: default_model_type(),\n ..Default::default()\n }\n }\n}\n\n#[derive(Debug, Clone, Serialize, Deserialize)]\npub struct ProviderModels {\n pub provider: String,\n pub models: Vec,\n}\n\nfn default_model_type() -> String {\n \"chat\".into()\n}\n\n#[derive(Debug, Clone, Copy, PartialEq)]\npub enum ModelType {\n Chat,\n Embedding,\n Reranker,\n}\n\nimpl Display for ModelType {\n fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {\n match self {\n ModelType::Chat => write!(f, \"chat\"),\n ModelType::Embedding => write!(f, \"embedding\"),\n ModelType::Reranker => write!(f, \"reranker\"),\n }\n }\n}\n\nimpl ModelType {\n pub fn can_create_from_name(self) -> bool {\n match self {\n ModelType::Chat => true,\n ModelType::Embedding => false,\n ModelType::Reranker => true,\n }\n }\n\n pub fn api_name(self) -> &'static str {\n match self {\n ModelType::Chat => \"chat_completions\",\n ModelType::Embedding => \"embeddings\",\n ModelType::Reranker => \"rerank\",\n }\n }\n\n pub fn extract_patch(self, patch: &RequestPatch) -> Option<&ApiPatch> {\n match self {\n ModelType::Chat => patch.chat_completions.as_ref(),\n ModelType::Embedding => patch.embeddings.as_ref(),\n ModelType::Reranker => patch.rerank.as_ref(),\n }\n }\n}\n\nfn stringify_option_value(value: &Option) -> String\nwhere\n T: std::fmt::Display,\n{\n match value {\n Some(value) => value.to_string(),\n None => \"-\".to_string(),\n }\n}\n" }, { "path": "src/client/openai.rs", "content": "use super::*;\n\nuse crate::utils::strip_think_tag;\n\nuse anyhow::{bail, Context, Result};\nuse reqwest::RequestBuilder;\nuse serde::Deserialize;\nuse serde_json::{json, Value};\n\nconst API_BASE: &str = \"https://api.openai.com/v1\";\n\n#[derive(Debug, Clone, Deserialize, Default)]\npub struct OpenAIConfig {\n pub name: Option,\n pub api_key: Option,\n pub api_base: Option,\n pub organization_id: Option,\n #[serde(default)]\n pub models: Vec,\n pub patch: Option,\n pub extra: Option,\n}\n\nimpl OpenAIClient {\n config_get_fn!(api_key, get_api_key);\n config_get_fn!(api_base, get_api_base);\n\n pub const PROMPTS: [PromptAction<'static>; 1] = [(\"api_key\", \"API Key\", None)];\n}\n\nimpl_client_trait!(\n OpenAIClient,\n (\n prepare_chat_completions,\n openai_chat_completions,\n openai_chat_completions_streaming\n ),\n (prepare_embeddings, openai_embeddings),\n (noop_prepare_rerank, noop_rerank),\n);\n\nfn prepare_chat_completions(\n self_: &OpenAIClient,\n data: ChatCompletionsData,\n) -> Result {\n let api_key = self_.get_api_key()?;\n let api_base = self_\n .get_api_base()\n .unwrap_or_else(|_| API_BASE.to_string());\n\n let url = format!(\"{}/chat/completions\", api_base.trim_end_matches('/'));\n\n let body = openai_build_chat_completions_body(data, &self_.model);\n\n let mut request_data = RequestData::new(url, body);\n\n request_data.bearer_auth(api_key);\n if let Some(organization_id) = &self_.config.organization_id {\n request_data.header(\"OpenAI-Organization\", organization_id);\n }\n\n Ok(request_data)\n}\n\nfn prepare_embeddings(self_: &OpenAIClient, data: &EmbeddingsData) -> Result {\n let api_key = self_.get_api_key()?;\n let api_base = self_\n .get_api_base()\n .unwrap_or_else(|_| API_BASE.to_string());\n\n let url = format!(\"{api_base}/embeddings\");\n\n let body = openai_build_embeddings_body(data, &self_.model);\n\n let mut request_data = RequestData::new(url, body);\n\n request_data.bearer_auth(api_key);\n if let Some(organization_id) = &self_.config.organization_id {\n request_data.header(\"OpenAI-Organization\", organization_id);\n }\n\n Ok(request_data)\n}\n\npub async fn openai_chat_completions(\n builder: RequestBuilder,\n _model: &Model,\n) -> Result {\n let res = builder.send().await?;\n let status = res.status();\n let data: Value = res.json().await?;\n if !status.is_success() {\n catch_error(&data, status.as_u16())?;\n }\n\n debug!(\"non-stream-data: {data}\");\n openai_extract_chat_completions(&data)\n}\n\npub async fn openai_chat_completions_streaming(\n builder: RequestBuilder,\n handler: &mut SseHandler,\n _model: &Model,\n) -> Result<()> {\n let mut call_id = String::new();\n let mut function_name = String::new();\n let mut function_arguments = String::new();\n let mut function_id = String::new();\n let mut reasoning_state = 0;\n let handle = |message: SseMmessage| -> Result {\n if message.data == \"[DONE]\" {\n if !function_name.is_empty() {\n if function_arguments.is_empty() {\n function_arguments = String::from(\"{}\");\n }\n let arguments: Value = function_arguments.parse().with_context(|| {\n format!(\"Tool call '{function_name}' have non-JSON arguments '{function_arguments}'\")\n })?;\n handler.tool_call(ToolCall::new(\n function_name.clone(),\n arguments,\n normalize_function_id(&function_id),\n ))?;\n }\n return Ok(true);\n }\n let data: Value = serde_json::from_str(&message.data)?;\n debug!(\"stream-data: {data}\");\n if let Some(text) = data[\"choices\"][0][\"delta\"][\"content\"]\n .as_str()\n .filter(|v| !v.is_empty())\n {\n if reasoning_state == 1 {\n handler.text(\"\\n\\n\\n\")?;\n reasoning_state = 0;\n }\n handler.text(text)?;\n } else if let Some(text) = data[\"choices\"][0][\"delta\"][\"reasoning_content\"]\n .as_str()\n .or_else(|| data[\"choices\"][0][\"delta\"][\"reasoning\"].as_str())\n .filter(|v| !v.is_empty())\n {\n if reasoning_state == 0 {\n handler.text(\"\\n\")?;\n reasoning_state = 1;\n }\n handler.text(text)?;\n }\n if let (Some(function), index, id) = (\n data[\"choices\"][0][\"delta\"][\"tool_calls\"][0][\"function\"].as_object(),\n data[\"choices\"][0][\"delta\"][\"tool_calls\"][0][\"index\"].as_u64(),\n data[\"choices\"][0][\"delta\"][\"tool_calls\"][0][\"id\"]\n .as_str()\n .filter(|v| !v.is_empty()),\n ) {\n if reasoning_state == 1 {\n handler.text(\"\\n\\n\\n\")?;\n reasoning_state = 0;\n }\n let maybe_call_id = format!(\"{}/{}\", id.unwrap_or_default(), index.unwrap_or_default());\n if maybe_call_id != call_id && maybe_call_id.len() >= call_id.len() {\n if !function_name.is_empty() {\n if function_arguments.is_empty() {\n function_arguments = String::from(\"{}\");\n }\n let arguments: Value = function_arguments.parse().with_context(|| {\n format!(\"Tool call '{function_name}' have non-JSON arguments '{function_arguments}'\")\n })?;\n handler.tool_call(ToolCall::new(\n function_name.clone(),\n arguments,\n normalize_function_id(&function_id),\n ))?;\n }\n function_name.clear();\n function_arguments.clear();\n function_id.clear();\n call_id = maybe_call_id;\n }\n if let Some(name) = function.get(\"name\").and_then(|v| v.as_str()) {\n if name.starts_with(&function_name) {\n function_name = name.to_string();\n } else {\n function_name.push_str(name);\n }\n }\n if let Some(arguments) = function.get(\"arguments\").and_then(|v| v.as_str()) {\n function_arguments.push_str(arguments);\n }\n if let Some(id) = id {\n function_id = id.to_string();\n }\n }\n Ok(false)\n };\n\n sse_stream(builder, handle).await\n}\n\npub async fn openai_embeddings(\n builder: RequestBuilder,\n _model: &Model,\n) -> Result {\n let res = builder.send().await?;\n let status = res.status();\n let data: Value = res.json().await?;\n if !status.is_success() {\n catch_error(&data, status.as_u16())?;\n }\n let res_body: EmbeddingsResBody =\n serde_json::from_value(data).context(\"Invalid embeddings data\")?;\n let output = res_body.data.into_iter().map(|v| v.embedding).collect();\n Ok(output)\n}\n\n#[derive(Deserialize)]\nstruct EmbeddingsResBody {\n data: Vec,\n}\n\n#[derive(Deserialize)]\nstruct EmbeddingsResBodyEmbedding {\n embedding: Vec,\n}\n\npub fn openai_build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Value {\n let ChatCompletionsData {\n messages,\n temperature,\n top_p,\n functions,\n stream,\n } = data;\n\n let messages_len = messages.len();\n let messages: Vec = messages\n .into_iter()\n .enumerate()\n .flat_map(|(i, message)| {\n let Message { role, content } = message;\n match content {\n MessageContent::ToolCalls(MessageContentToolCalls {\n tool_results,\n text: _,\n sequence,\n }) => {\n if !sequence {\n let tool_calls: Vec<_> = tool_results\n .iter()\n .map(|tool_result| {\n json!({\n \"id\": tool_result.call.id,\n \"type\": \"function\",\n \"function\": {\n \"name\": tool_result.call.name,\n \"arguments\": tool_result.call.arguments.to_string(),\n },\n })\n })\n .collect();\n let mut messages = vec![\n json!({ \"role\": MessageRole::Assistant, \"tool_calls\": tool_calls }),\n ];\n for tool_result in tool_results {\n messages.push(json!({\n \"role\": \"tool\",\n \"content\": tool_result.output.to_string(),\n \"tool_call_id\": tool_result.call.id,\n }));\n }\n messages\n } else {\n tool_results.into_iter().flat_map(|tool_result| {\n vec![\n json!({\n \"role\": MessageRole::Assistant,\n \"tool_calls\": [\n {\n \"id\": tool_result.call.id,\n \"type\": \"function\",\n \"function\": {\n \"name\": tool_result.call.name,\n \"arguments\": tool_result.call.arguments.to_string(),\n },\n }\n ]\n }),\n json!({\n \"role\": \"tool\",\n \"content\": tool_result.output.to_string(),\n \"tool_call_id\": tool_result.call.id,\n })\n ]\n\n }).collect()\n }\n }\n MessageContent::Text(text) if role.is_assistant() && i != messages_len - 1 => {\n vec![json!({ \"role\": role, \"content\": strip_think_tag(&text) }\n )]\n }\n _ => vec![json!({ \"role\": role, \"content\": content })],\n }\n })\n .collect();\n\n let mut body = json!({\n \"model\": &model.real_name(),\n \"messages\": messages,\n });\n\n if let Some(v) = model.max_tokens_param() {\n if model\n .patch()\n .and_then(|v| v.get(\"body\").and_then(|v| v.get(\"max_tokens\")))\n == Some(&Value::Null)\n {\n body[\"max_completion_tokens\"] = v.into();\n } else {\n body[\"max_tokens\"] = v.into();\n }\n }\n if let Some(v) = temperature {\n body[\"temperature\"] = v.into();\n }\n if let Some(v) = top_p {\n body[\"top_p\"] = v.into();\n }\n if stream {\n body[\"stream\"] = true.into();\n }\n if let Some(functions) = functions {\n body[\"tools\"] = functions\n .iter()\n .map(|v| {\n json!({\n \"type\": \"function\",\n \"function\": v,\n })\n })\n .collect();\n }\n body\n}\n\npub fn openai_build_embeddings_body(data: &EmbeddingsData, model: &Model) -> Value {\n json!({\n \"input\": data.texts,\n \"model\": model.real_name()\n })\n}\n\npub fn openai_extract_chat_completions(data: &Value) -> Result {\n let text = data[\"choices\"][0][\"message\"][\"content\"]\n .as_str()\n .unwrap_or_default();\n\n let reasoning = data[\"choices\"][0][\"message\"][\"reasoning_content\"]\n .as_str()\n .or_else(|| data[\"choices\"][0][\"message\"][\"reasoning\"].as_str())\n .unwrap_or_default()\n .trim();\n\n let mut tool_calls = vec![];\n if let Some(calls) = data[\"choices\"][0][\"message\"][\"tool_calls\"].as_array() {\n for call in calls {\n if let (Some(name), Some(arguments), Some(id)) = (\n call[\"function\"][\"name\"].as_str(),\n call[\"function\"][\"arguments\"].as_str(),\n call[\"id\"].as_str(),\n ) {\n let arguments: Value = arguments.parse().with_context(|| {\n format!(\"Tool call '{name}' have non-JSON arguments '{arguments}'\")\n })?;\n tool_calls.push(ToolCall::new(\n name.to_string(),\n arguments,\n Some(id.to_string()),\n ));\n }\n }\n };\n\n if text.is_empty() && tool_calls.is_empty() {\n bail!(\"Invalid response data: {data}\");\n }\n let text = if !reasoning.is_empty() {\n format!(\"\\n{reasoning}\\n\\n\\n{text}\")\n } else {\n text.to_string()\n };\n let output = ChatCompletionsOutput {\n text,\n tool_calls,\n id: data[\"id\"].as_str().map(|v| v.to_string()),\n input_tokens: data[\"usage\"][\"prompt_tokens\"].as_u64(),\n output_tokens: data[\"usage\"][\"completion_tokens\"].as_u64(),\n };\n Ok(output)\n}\n\nfn normalize_function_id(value: &str) -> Option {\n if value.is_empty() {\n None\n } else {\n Some(value.to_string())\n }\n}\n" }, { "path": "src/client/openai_compatible.rs", "content": "use super::openai::*;\nuse super::*;\n\nuse anyhow::{Context, Result};\nuse reqwest::RequestBuilder;\nuse serde::Deserialize;\nuse serde_json::{json, Value};\n\n#[derive(Debug, Clone, Deserialize)]\npub struct OpenAICompatibleConfig {\n pub name: Option,\n pub api_base: Option,\n pub api_key: Option,\n #[serde(default)]\n pub models: Vec,\n pub patch: Option,\n pub extra: Option,\n}\n\nimpl OpenAICompatibleClient {\n config_get_fn!(api_base, get_api_base);\n config_get_fn!(api_key, get_api_key);\n\n pub const PROMPTS: [PromptAction<'static>; 0] = [];\n}\n\nimpl_client_trait!(\n OpenAICompatibleClient,\n (\n prepare_chat_completions,\n openai_chat_completions,\n openai_chat_completions_streaming\n ),\n (prepare_embeddings, openai_embeddings),\n (prepare_rerank, generic_rerank),\n);\n\nfn prepare_chat_completions(\n self_: &OpenAICompatibleClient,\n data: ChatCompletionsData,\n) -> Result {\n let api_key = self_.get_api_key().ok();\n let api_base = get_api_base_ext(self_)?;\n\n let url = format!(\"{api_base}/chat/completions\");\n\n let body = openai_build_chat_completions_body(data, &self_.model);\n\n let mut request_data = RequestData::new(url, body);\n\n if let Some(api_key) = api_key {\n request_data.bearer_auth(api_key);\n }\n\n Ok(request_data)\n}\n\nfn prepare_embeddings(\n self_: &OpenAICompatibleClient,\n data: &EmbeddingsData,\n) -> Result {\n let api_key = self_.get_api_key().ok();\n let api_base = get_api_base_ext(self_)?;\n\n let url = format!(\"{api_base}/embeddings\");\n\n let body = openai_build_embeddings_body(data, &self_.model);\n\n let mut request_data = RequestData::new(url, body);\n\n if let Some(api_key) = api_key {\n request_data.bearer_auth(api_key);\n }\n\n Ok(request_data)\n}\n\nfn prepare_rerank(self_: &OpenAICompatibleClient, data: &RerankData) -> Result {\n let api_key = self_.get_api_key().ok();\n let api_base = get_api_base_ext(self_)?;\n\n let url = if self_.name().starts_with(\"ernie\") {\n format!(\"{api_base}/rerankers\")\n } else {\n format!(\"{api_base}/rerank\")\n };\n\n let body = generic_build_rerank_body(data, &self_.model);\n\n let mut request_data = RequestData::new(url, body);\n\n if let Some(api_key) = api_key {\n request_data.bearer_auth(api_key);\n }\n\n Ok(request_data)\n}\n\nfn get_api_base_ext(self_: &OpenAICompatibleClient) -> Result {\n let api_base = match self_.get_api_base() {\n Ok(v) => v,\n Err(err) => {\n match OPENAI_COMPATIBLE_PROVIDERS\n .into_iter()\n .find_map(|(name, api_base)| {\n if name == self_.model.client_name() {\n Some(api_base.to_string())\n } else {\n None\n }\n }) {\n Some(v) => v,\n None => return Err(err),\n }\n }\n };\n Ok(api_base.trim_end_matches('/').to_string())\n}\n\npub async fn generic_rerank(builder: RequestBuilder, _model: &Model) -> Result {\n let res = builder.send().await?;\n let status = res.status();\n let mut data: Value = res.json().await?;\n if !status.is_success() {\n catch_error(&data, status.as_u16())?;\n }\n if data.get(\"results\").is_none() && data.get(\"data\").is_some() {\n if let Some(data_obj) = data.as_object_mut() {\n if let Some(value) = data_obj.remove(\"data\") {\n data_obj.insert(\"results\".to_string(), value);\n }\n }\n }\n let res_body: GenericRerankResBody =\n serde_json::from_value(data).context(\"Invalid rerank data\")?;\n Ok(res_body.results)\n}\n\n#[derive(Deserialize)]\npub struct GenericRerankResBody {\n pub results: RerankOutput,\n}\n\npub fn generic_build_rerank_body(data: &RerankData, model: &Model) -> Value {\n let RerankData {\n query,\n documents,\n top_n,\n } = data;\n\n let mut body = json!({\n \"model\": model.real_name(),\n \"query\": query,\n \"documents\": documents,\n });\n if model.client_name().starts_with(\"voyageai\") {\n body[\"top_k\"] = (*top_n).into()\n } else {\n body[\"top_n\"] = (*top_n).into()\n }\n body\n}\n" }, { "path": "src/client/stream.rs", "content": "use super::{catch_error, ToolCall};\nuse crate::utils::AbortSignal;\n\nuse anyhow::{anyhow, bail, Context, Result};\nuse futures_util::{Stream, StreamExt};\nuse reqwest::RequestBuilder;\nuse reqwest_eventsource::{Error as EventSourceError, Event, RequestBuilderExt};\nuse serde_json::Value;\nuse tokio::sync::mpsc::UnboundedSender;\n\npub struct SseHandler {\n sender: UnboundedSender,\n abort_signal: AbortSignal,\n buffer: String,\n tool_calls: Vec,\n}\n\nimpl SseHandler {\n pub fn new(sender: UnboundedSender, abort_signal: AbortSignal) -> Self {\n Self {\n sender,\n abort_signal,\n buffer: String::new(),\n tool_calls: Vec::new(),\n }\n }\n\n pub fn text(&mut self, text: &str) -> Result<()> {\n // debug!(\"HandleText: {}\", text);\n if text.is_empty() {\n return Ok(());\n }\n self.buffer.push_str(text);\n let ret = self\n .sender\n .send(SseEvent::Text(text.to_string()))\n .with_context(|| \"Failed to send SseEvent:Text\");\n if let Err(err) = ret {\n if self.abort_signal.aborted() {\n return Ok(());\n }\n return Err(err);\n }\n Ok(())\n }\n\n pub fn done(&mut self) {\n // debug!(\"HandleDone\");\n let ret = self.sender.send(SseEvent::Done);\n if ret.is_err() {\n if self.abort_signal.aborted() {\n return;\n }\n warn!(\"Failed to send SseEvent:Done\");\n }\n }\n\n pub fn tool_call(&mut self, call: ToolCall) -> Result<()> {\n // debug!(\"HandleCall: {:?}\", call);\n self.tool_calls.push(call);\n Ok(())\n }\n\n pub fn abort(&self) -> AbortSignal {\n self.abort_signal.clone()\n }\n\n pub fn tool_calls(&self) -> &[ToolCall] {\n &self.tool_calls\n }\n\n pub fn take(self) -> (String, Vec) {\n let Self {\n buffer, tool_calls, ..\n } = self;\n (buffer, tool_calls)\n }\n}\n\n#[derive(Debug)]\npub enum SseEvent {\n Text(String),\n Done,\n}\n\n#[derive(Debug)]\npub struct SseMmessage {\n #[allow(unused)]\n pub event: String,\n pub data: String,\n}\n\npub async fn sse_stream(builder: RequestBuilder, mut handle: F) -> Result<()>\nwhere\n F: FnMut(SseMmessage) -> Result,\n{\n let mut es = builder.eventsource()?;\n while let Some(event) = es.next().await {\n match event {\n Ok(Event::Open) => {}\n Ok(Event::Message(message)) => {\n let message = SseMmessage {\n event: message.event,\n data: message.data,\n };\n if handle(message)? {\n break;\n }\n }\n Err(err) => {\n match err {\n EventSourceError::StreamEnded => {}\n EventSourceError::InvalidStatusCode(status, res) => {\n let text = res.text().await?;\n let data: Value = match text.parse() {\n Ok(data) => data,\n Err(_) => {\n bail!(\n \"Invalid response data: {text} (status: {})\",\n status.as_u16()\n );\n }\n };\n catch_error(&data, status.as_u16())?;\n }\n EventSourceError::InvalidContentType(header_value, res) => {\n let text = res.text().await?;\n bail!(\n \"Invalid response event-stream. content-type: {}, data: {text}\",\n header_value.to_str().unwrap_or_default()\n );\n }\n _ => {\n bail!(\"{}\", err);\n }\n }\n es.close();\n }\n }\n }\n Ok(())\n}\n\npub async fn json_stream(mut stream: S, mut handle: F) -> Result<()>\nwhere\n S: Stream> + Unpin,\n F: FnMut(&str) -> Result<()>,\n E: std::error::Error,\n{\n let mut parser = JsonStreamParser::default();\n let mut unparsed_bytes = vec![];\n while let Some(chunk_bytes) = stream.next().await {\n let chunk_bytes =\n chunk_bytes.map_err(|err| anyhow!(\"Failed to read json stream, {err}\"))?;\n unparsed_bytes.extend(chunk_bytes);\n match std::str::from_utf8(&unparsed_bytes) {\n Ok(text) => {\n parser.process(text, &mut handle)?;\n unparsed_bytes.clear();\n }\n Err(_) => {\n continue;\n }\n }\n }\n if !unparsed_bytes.is_empty() {\n let text = std::str::from_utf8(&unparsed_bytes)?;\n parser.process(text, &mut handle)?;\n }\n\n Ok(())\n}\n\n#[derive(Debug, Default)]\nstruct JsonStreamParser {\n buffer: Vec,\n cursor: usize,\n start: Option,\n balances: Vec,\n quoting: bool,\n escape: bool,\n}\n\nimpl JsonStreamParser {\n fn process(&mut self, text: &str, handle: &mut F) -> Result<()>\n where\n F: FnMut(&str) -> Result<()>,\n {\n self.buffer.extend(text.chars());\n\n for i in self.cursor..self.buffer.len() {\n let ch = self.buffer[i];\n if self.quoting {\n if ch == '\\\\' {\n self.escape = !self.escape;\n } else {\n if !self.escape && ch == '\"' {\n self.quoting = false;\n }\n self.escape = false;\n }\n continue;\n }\n match ch {\n '\"' => {\n self.quoting = true;\n self.escape = false;\n }\n '{' => {\n if self.balances.is_empty() {\n self.start = Some(i);\n }\n self.balances.push(ch);\n }\n '[' => {\n if self.start.is_some() {\n self.balances.push(ch);\n }\n }\n '}' => {\n self.balances.pop();\n if self.balances.is_empty() {\n if let Some(start) = self.start.take() {\n let value: String = self.buffer[start..=i].iter().collect();\n handle(&value)?;\n }\n }\n }\n ']' => {\n self.balances.pop();\n }\n _ => {}\n }\n }\n self.cursor = self.buffer.len();\n Ok(())\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n use bytes::Bytes;\n use futures_util::stream;\n use rand::Rng;\n\n fn split_chunks(text: &str) -> Vec> {\n let mut rng = rand::rng();\n let len = text.len();\n let cut1 = rng.random_range(1..len - 1);\n let cut2 = rng.random_range(cut1 + 1..len);\n let chunk1 = text.as_bytes()[..cut1].to_vec();\n let chunk2 = text.as_bytes()[cut1..cut2].to_vec();\n let chunk3 = text.as_bytes()[cut2..].to_vec();\n vec![chunk1, chunk2, chunk3]\n }\n\n macro_rules! assert_json_stream {\n ($input:expr, $output:expr) => {\n let chunks: Vec<_> = split_chunks($input)\n .into_iter()\n .map(|chunk| Ok::<_, std::convert::Infallible>(Bytes::from(chunk)))\n .collect();\n let stream = stream::iter(chunks);\n let mut output = vec![];\n let ret = json_stream(stream, |data| {\n output.push(data.to_string());\n Ok(())\n })\n .await;\n assert!(ret.is_ok());\n assert_eq!($output.replace(\"\\r\\n\", \"\\n\"), output.join(\"\\n\"))\n };\n }\n\n #[tokio::test]\n async fn test_json_stream_ndjson() {\n let data = r#\"{\"key\": \"value\"}\n{\"key\": \"value2\"}\n{\"key\": \"value3\"}\"#;\n assert_json_stream!(data, data);\n }\n\n #[tokio::test]\n async fn test_json_stream_array() {\n let input = r#\"[\n{\"key\": \"value\"},\n{\"key\": \"value2\"},\n{\"key\": \"value3\"},\"#;\n let output = r#\"{\"key\": \"value\"}\n{\"key\": \"value2\"}\n{\"key\": \"value3\"}\"#;\n assert_json_stream!(input, output);\n }\n}\n" }, { "path": "src/client/vertexai.rs", "content": "use super::access_token::*;\nuse super::claude::*;\nuse super::openai::*;\nuse super::*;\n\nuse anyhow::{anyhow, bail, Context, Result};\nuse chrono::{Duration, Utc};\nuse reqwest::{Client as ReqwestClient, RequestBuilder};\nuse serde::Deserialize;\nuse serde_json::{json, Value};\nuse std::{path::PathBuf, str::FromStr};\n\n#[derive(Debug, Clone, Deserialize, Default)]\npub struct VertexAIConfig {\n pub name: Option,\n pub project_id: Option,\n pub location: Option,\n pub adc_file: Option,\n #[serde(default)]\n pub models: Vec,\n pub patch: Option,\n pub extra: Option,\n}\n\nimpl VertexAIClient {\n config_get_fn!(project_id, get_project_id);\n config_get_fn!(location, get_location);\n\n pub const PROMPTS: [PromptAction<'static>; 2] = [\n (\"project_id\", \"Project ID\", None),\n (\"location\", \"Location\", None),\n ];\n}\n\n#[async_trait::async_trait]\nimpl Client for VertexAIClient {\n client_common_fns!();\n\n async fn chat_completions_inner(\n &self,\n client: &ReqwestClient,\n data: ChatCompletionsData,\n ) -> Result {\n prepare_gcloud_access_token(client, self.name(), &self.config.adc_file).await?;\n let model = self.model();\n let model_category = ModelCategory::from_str(model.real_name())?;\n let request_data = prepare_chat_completions(self, data, &model_category)?;\n let builder = self.request_builder(client, request_data);\n match model_category {\n ModelCategory::Gemini => gemini_chat_completions(builder, model).await,\n ModelCategory::Claude => claude_chat_completions(builder, model).await,\n ModelCategory::Mistral => openai_chat_completions(builder, model).await,\n }\n }\n\n async fn chat_completions_streaming_inner(\n &self,\n client: &ReqwestClient,\n handler: &mut SseHandler,\n data: ChatCompletionsData,\n ) -> Result<()> {\n prepare_gcloud_access_token(client, self.name(), &self.config.adc_file).await?;\n let model = self.model();\n let model_category = ModelCategory::from_str(model.real_name())?;\n let request_data = prepare_chat_completions(self, data, &model_category)?;\n let builder = self.request_builder(client, request_data);\n match model_category {\n ModelCategory::Gemini => {\n gemini_chat_completions_streaming(builder, handler, model).await\n }\n ModelCategory::Claude => {\n claude_chat_completions_streaming(builder, handler, model).await\n }\n ModelCategory::Mistral => {\n openai_chat_completions_streaming(builder, handler, model).await\n }\n }\n }\n\n async fn embeddings_inner(\n &self,\n client: &ReqwestClient,\n data: &EmbeddingsData,\n ) -> Result>> {\n prepare_gcloud_access_token(client, self.name(), &self.config.adc_file).await?;\n let request_data = prepare_embeddings(self, data)?;\n let builder = self.request_builder(client, request_data);\n embeddings(builder, self.model()).await\n }\n}\n\nfn prepare_chat_completions(\n self_: &VertexAIClient,\n data: ChatCompletionsData,\n model_category: &ModelCategory,\n) -> Result {\n let project_id = self_.get_project_id()?;\n let location = self_.get_location()?;\n let access_token = get_access_token(self_.name())?;\n\n let base_url = if location == \"global\" {\n format!(\"https://aiplatform.googleapis.com/v1/projects/{project_id}/locations/global/publishers\")\n } else {\n format!(\"https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/publishers\")\n };\n\n let model_name = self_.model.real_name();\n\n let url = match model_category {\n ModelCategory::Gemini => {\n let func = match data.stream {\n true => \"streamGenerateContent\",\n false => \"generateContent\",\n };\n format!(\"{base_url}/google/models/{model_name}:{func}\")\n }\n ModelCategory::Claude => {\n format!(\"{base_url}/anthropic/models/{model_name}:streamRawPredict\")\n }\n ModelCategory::Mistral => {\n let func = match data.stream {\n true => \"streamRawPredict\",\n false => \"rawPredict\",\n };\n format!(\"{base_url}/mistralai/models/{model_name}:{func}\")\n }\n };\n\n let body = match model_category {\n ModelCategory::Gemini => gemini_build_chat_completions_body(data, &self_.model)?,\n ModelCategory::Claude => {\n let mut body = claude_build_chat_completions_body(data, &self_.model)?;\n if let Some(body_obj) = body.as_object_mut() {\n body_obj.remove(\"model\");\n }\n body[\"anthropic_version\"] = \"vertex-2023-10-16\".into();\n body\n }\n ModelCategory::Mistral => {\n let mut body = openai_build_chat_completions_body(data, &self_.model);\n if let Some(body_obj) = body.as_object_mut() {\n body_obj[\"model\"] = strip_model_version(self_.model.real_name()).into();\n }\n body\n }\n };\n\n let mut request_data = RequestData::new(url, body);\n\n request_data.bearer_auth(access_token);\n\n Ok(request_data)\n}\n\nfn prepare_embeddings(self_: &VertexAIClient, data: &EmbeddingsData) -> Result {\n let project_id = self_.get_project_id()?;\n let location = self_.get_location()?;\n let access_token = get_access_token(self_.name())?;\n\n let base_url = if location == \"global\" {\n format!(\"https://aiplatform.googleapis.com/v1/projects/{project_id}/locations/global/publishers\")\n } else {\n format!(\"https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/publishers\")\n };\n let url = format!(\n \"{base_url}/google/models/{}:predict\",\n self_.model.real_name()\n );\n\n let instances: Vec<_> = data.texts.iter().map(|v| json!({\"content\": v})).collect();\n\n let body = json!({\n \"instances\": instances,\n });\n\n let mut request_data = RequestData::new(url, body);\n\n request_data.bearer_auth(access_token);\n\n Ok(request_data)\n}\n\npub async fn gemini_chat_completions(\n builder: RequestBuilder,\n _model: &Model,\n) -> Result {\n let res = builder.send().await?;\n let status = res.status();\n let data: Value = res.json().await?;\n if !status.is_success() {\n catch_error(&data, status.as_u16())?;\n }\n debug!(\"non-stream-data: {data}\");\n gemini_extract_chat_completions_text(&data)\n}\n\npub async fn gemini_chat_completions_streaming(\n builder: RequestBuilder,\n handler: &mut SseHandler,\n _model: &Model,\n) -> Result<()> {\n let res = builder.send().await?;\n let status = res.status();\n if !status.is_success() {\n let data: Value = res.json().await?;\n catch_error(&data, status.as_u16())?;\n } else {\n let handle = |value: &str| -> Result<()> {\n let data: Value = serde_json::from_str(value)?;\n debug!(\"stream-data: {data}\");\n if let Some(parts) = data[\"candidates\"][0][\"content\"][\"parts\"].as_array() {\n for (i, part) in parts.iter().enumerate() {\n if let Some(text) = part[\"text\"].as_str() {\n if i > 0 {\n handler.text(\"\\n\\n\")?;\n }\n handler.text(text)?;\n } else if let (Some(name), Some(args)) = (\n part[\"functionCall\"][\"name\"].as_str(),\n part[\"functionCall\"][\"args\"].as_object(),\n ) {\n handler.tool_call(ToolCall::new(name.to_string(), json!(args), None))?;\n }\n }\n } else if let Some(\"SAFETY\") = data[\"promptFeedback\"][\"blockReason\"]\n .as_str()\n .or_else(|| data[\"candidates\"][0][\"finishReason\"].as_str())\n {\n bail!(\"Blocked due to safety\")\n }\n\n Ok(())\n };\n json_stream(res.bytes_stream(), handle).await?;\n }\n Ok(())\n}\n\nasync fn embeddings(builder: RequestBuilder, _model: &Model) -> Result {\n let res = builder.send().await?;\n let status = res.status();\n let data: Value = res.json().await?;\n if !status.is_success() {\n catch_error(&data, status.as_u16())?;\n }\n let res_body: EmbeddingsResBody =\n serde_json::from_value(data).context(\"Invalid embeddings data\")?;\n let output = res_body\n .predictions\n .into_iter()\n .map(|v| v.embeddings.values)\n .collect();\n Ok(output)\n}\n\n#[derive(Deserialize)]\nstruct EmbeddingsResBody {\n predictions: Vec,\n}\n\n#[derive(Deserialize)]\nstruct EmbeddingsResBodyPrediction {\n embeddings: EmbeddingsResBodyPredictionEmbeddings,\n}\n\n#[derive(Deserialize)]\nstruct EmbeddingsResBodyPredictionEmbeddings {\n values: Vec,\n}\n\nfn gemini_extract_chat_completions_text(data: &Value) -> Result {\n let mut text_parts = vec![];\n let mut tool_calls = vec![];\n if let Some(parts) = data[\"candidates\"][0][\"content\"][\"parts\"].as_array() {\n for part in parts {\n if let Some(text) = part[\"text\"].as_str() {\n text_parts.push(text);\n }\n if let (Some(name), Some(args)) = (\n part[\"functionCall\"][\"name\"].as_str(),\n part[\"functionCall\"][\"args\"].as_object(),\n ) {\n tool_calls.push(ToolCall::new(name.to_string(), json!(args), None));\n }\n }\n }\n\n let text = text_parts.join(\"\\n\\n\");\n if text.is_empty() && tool_calls.is_empty() {\n if let Some(\"SAFETY\") = data[\"promptFeedback\"][\"blockReason\"]\n .as_str()\n .or_else(|| data[\"candidates\"][0][\"finishReason\"].as_str())\n {\n bail!(\"Blocked due to safety\")\n } else {\n bail!(\"Invalid response data: {data}\");\n }\n }\n let output = ChatCompletionsOutput {\n text,\n tool_calls,\n id: None,\n input_tokens: data[\"usageMetadata\"][\"promptTokenCount\"].as_u64(),\n output_tokens: data[\"usageMetadata\"][\"candidatesTokenCount\"].as_u64(),\n };\n Ok(output)\n}\n\npub fn gemini_build_chat_completions_body(\n data: ChatCompletionsData,\n model: &Model,\n) -> Result {\n let ChatCompletionsData {\n mut messages,\n temperature,\n top_p,\n functions,\n stream: _,\n } = data;\n\n let system_message = extract_system_message(&mut messages);\n\n let mut network_image_urls = vec![];\n let contents: Vec = messages\n .into_iter()\n .flat_map(|message| {\n let Message { role, content } = message;\n let role = match role {\n MessageRole::User => \"user\",\n _ => \"model\",\n };\n match content {\n MessageContent::Text(text) => vec![json!({\n \"role\": role,\n \"parts\": [{ \"text\": text }]\n })],\n MessageContent::Array(list) => {\n let parts: Vec = list\n .into_iter()\n .map(|item| match item {\n MessageContentPart::Text { text } => json!({\"text\": text}),\n MessageContentPart::ImageUrl { image_url: ImageUrl { url } } => {\n if let Some((mime_type, data)) = url.strip_prefix(\"data:\").and_then(|v| v.split_once(\";base64,\")) {\n json!({ \"inline_data\": { \"mime_type\": mime_type, \"data\": data } })\n } else {\n network_image_urls.push(url.clone());\n json!({ \"url\": url })\n }\n },\n })\n .collect();\n vec![json!({ \"role\": role, \"parts\": parts })]\n },\n MessageContent::ToolCalls(MessageContentToolCalls { tool_results, .. }) => {\n let model_parts: Vec = tool_results.iter().map(|tool_result| {\n json!({\n \"functionCall\": {\n \"name\": tool_result.call.name,\n \"args\": tool_result.call.arguments,\n }\n })\n }).collect();\n let function_parts: Vec = tool_results.into_iter().map(|tool_result| {\n json!({\n \"functionResponse\": {\n \"name\": tool_result.call.name,\n \"response\": {\n \"name\": tool_result.call.name,\n \"content\": tool_result.output,\n }\n }\n })\n }).collect();\n vec![\n json!({ \"role\": \"model\", \"parts\": model_parts }),\n json!({ \"role\": \"function\", \"parts\": function_parts }),\n ]\n }\n }\n })\n .collect();\n\n if !network_image_urls.is_empty() {\n bail!(\n \"The model does not support network images: {:?}\",\n network_image_urls\n );\n }\n\n let mut body = json!({ \"contents\": contents, \"generationConfig\": {} });\n\n if let Some(v) = system_message {\n body[\"systemInstruction\"] = json!({ \"parts\": [{\"text\": v }] });\n }\n\n if let Some(v) = model.max_tokens_param() {\n body[\"generationConfig\"][\"maxOutputTokens\"] = v.into();\n }\n if let Some(v) = temperature {\n body[\"generationConfig\"][\"temperature\"] = v.into();\n }\n if let Some(v) = top_p {\n body[\"generationConfig\"][\"topP\"] = v.into();\n }\n\n if let Some(functions) = functions {\n // Gemini doesn't support functions with parameters that have empty properties, so we need to patch it.\n let function_declarations: Vec<_> = functions\n .into_iter()\n .map(|function| {\n if function.parameters.is_empty_properties() {\n json!({\n \"name\": function.name,\n \"description\": function.description,\n })\n } else {\n json!(function)\n }\n })\n .collect();\n body[\"tools\"] = json!([{ \"functionDeclarations\": function_declarations }]);\n }\n\n Ok(body)\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\nenum ModelCategory {\n Gemini,\n Claude,\n Mistral,\n}\n\nimpl FromStr for ModelCategory {\n type Err = anyhow::Error;\n\n fn from_str(s: &str) -> std::result::Result {\n if s.starts_with(\"gemini\") {\n Ok(ModelCategory::Gemini)\n } else if s.starts_with(\"claude\") {\n Ok(ModelCategory::Claude)\n } else if s.starts_with(\"mistral\") || s.starts_with(\"codestral\") {\n Ok(ModelCategory::Mistral)\n } else {\n unsupported_model!(s)\n }\n }\n}\n\npub async fn prepare_gcloud_access_token(\n client: &reqwest::Client,\n client_name: &str,\n adc_file: &Option,\n) -> Result<()> {\n if !is_valid_access_token(client_name) {\n let (token, expires_in) = fetch_access_token(client, adc_file)\n .await\n .with_context(|| \"Failed to fetch access token\")?;\n let expires_at = Utc::now()\n + Duration::try_seconds(expires_in)\n .ok_or_else(|| anyhow!(\"Failed to parse expires_in of access_token\"))?;\n set_access_token(client_name, token, expires_at.timestamp())\n }\n Ok(())\n}\n\nasync fn fetch_access_token(\n client: &reqwest::Client,\n file: &Option,\n) -> Result<(String, i64)> {\n let credentials = load_adc(file).await?;\n let value: Value = client\n .post(\"https://oauth2.googleapis.com/token\")\n .json(&credentials)\n .send()\n .await?\n .json()\n .await?;\n\n if let (Some(access_token), Some(expires_in)) =\n (value[\"access_token\"].as_str(), value[\"expires_in\"].as_i64())\n {\n Ok((access_token.to_string(), expires_in))\n } else if let Some(err_msg) = value[\"error_description\"].as_str() {\n bail!(\"{err_msg}\")\n } else {\n bail!(\"Invalid response data: {value}\")\n }\n}\n\nasync fn load_adc(file: &Option) -> Result {\n let adc_file = file\n .as_ref()\n .map(PathBuf::from)\n .or_else(default_adc_file)\n .ok_or_else(|| anyhow!(\"No application_default_credentials.json\"))?;\n let data = tokio::fs::read_to_string(adc_file).await?;\n let data: Value = serde_json::from_str(&data)?;\n if let (Some(client_id), Some(client_secret), Some(refresh_token)) = (\n data[\"client_id\"].as_str(),\n data[\"client_secret\"].as_str(),\n data[\"refresh_token\"].as_str(),\n ) {\n Ok(json!({\n \"client_id\": client_id,\n \"client_secret\": client_secret,\n \"refresh_token\": refresh_token,\n \"grant_type\": \"refresh_token\",\n }))\n } else {\n bail!(\"Invalid application_default_credentials.json\")\n }\n}\n\n#[cfg(not(windows))]\nfn default_adc_file() -> Option {\n let mut path = dirs::home_dir()?;\n path.push(\".config\");\n path.push(\"gcloud\");\n path.push(\"application_default_credentials.json\");\n Some(path)\n}\n\n#[cfg(windows)]\nfn default_adc_file() -> Option {\n let mut path = dirs::config_dir()?;\n path.push(\"gcloud\");\n path.push(\"application_default_credentials.json\");\n Some(path)\n}\n\nfn strip_model_version(name: &str) -> &str {\n match name.split_once('@') {\n Some((v, _)) => v,\n None => name,\n }\n}\n" }, { "path": "src/config/agent.rs", "content": "use super::*;\n\nuse crate::{\n client::Model,\n function::{run_llm_function, Functions},\n};\n\nuse anyhow::{Context, Result};\nuse inquire::{validator::Validation, Text};\nuse std::{fs::read_to_string, path::Path};\n\nuse serde::{Deserialize, Serialize};\n\nconst DEFAULT_AGENT_NAME: &str = \"rag\";\n\npub type AgentVariables = IndexMap;\n\n#[derive(Debug, Clone)]\npub struct Agent {\n name: String,\n config: AgentConfig,\n definition: AgentDefinition,\n shared_variables: AgentVariables,\n session_variables: Option,\n shared_dynamic_instructions: Option,\n session_dynamic_instructions: Option,\n functions: Functions,\n rag: Option>,\n model: Model,\n}\n\nimpl Agent {\n pub async fn init(\n config: &GlobalConfig,\n name: &str,\n abort_signal: AbortSignal,\n ) -> Result {\n let functions_dir = Config::agent_functions_dir(name);\n let definition_file_path = functions_dir.join(\"index.yaml\");\n if !definition_file_path.exists() {\n bail!(\"Unknown agent `{name}`\");\n }\n let functions_file_path = functions_dir.join(\"functions.json\");\n let rag_path = Config::agent_rag_file(name, DEFAULT_AGENT_NAME);\n let config_path = Config::agent_config_file(name);\n let mut agent_config = if config_path.exists() {\n AgentConfig::load(&config_path)?\n } else {\n AgentConfig::new(&config.read())\n };\n let mut definition = AgentDefinition::load(&definition_file_path)?;\n let functions = if functions_file_path.exists() {\n Functions::init(&functions_file_path)?\n } else {\n Functions::default()\n };\n definition.replace_tools_placeholder(&functions);\n\n agent_config.load_envs(&definition.name);\n\n let model = {\n let config = config.read();\n match agent_config.model_id.as_ref() {\n Some(model_id) => Model::retrieve_model(&config, model_id, ModelType::Chat)?,\n None => {\n if agent_config.temperature.is_none() {\n agent_config.temperature = config.temperature;\n }\n if agent_config.top_p.is_none() {\n agent_config.top_p = config.top_p;\n }\n config.current_model().clone()\n }\n }\n };\n\n let rag = if rag_path.exists() {\n Some(Arc::new(Rag::load(config, DEFAULT_AGENT_NAME, &rag_path)?))\n } else if !definition.documents.is_empty() && !config.read().info_flag {\n let mut ans = false;\n if *IS_STDOUT_TERMINAL {\n ans = Confirm::new(\"The agent has the documents, init RAG?\")\n .with_default(true)\n .prompt()?;\n }\n if ans {\n let mut document_paths = vec![];\n for path in &definition.documents {\n if is_url(path) {\n document_paths.push(path.to_string());\n } else {\n let new_path = safe_join_path(&functions_dir, path)\n .ok_or_else(|| anyhow!(\"Invalid document path: '{path}'\"))?;\n document_paths.push(new_path.display().to_string())\n }\n }\n let rag =\n Rag::init(config, \"rag\", &rag_path, &document_paths, abort_signal).await?;\n Some(Arc::new(rag))\n } else {\n None\n }\n } else {\n None\n };\n\n Ok(Self {\n name: name.to_string(),\n config: agent_config,\n definition,\n shared_variables: Default::default(),\n session_variables: None,\n shared_dynamic_instructions: None,\n session_dynamic_instructions: None,\n functions,\n rag,\n model,\n })\n }\n\n pub fn init_agent_variables(\n agent_variables: &[AgentVariable],\n variables: &AgentVariables,\n no_interaction: bool,\n ) -> Result {\n let mut output = IndexMap::new();\n if agent_variables.is_empty() {\n return Ok(output);\n }\n let mut printed = false;\n let mut unset_variables = vec![];\n for agent_variable in agent_variables {\n let key = agent_variable.name.clone();\n match variables.get(&key) {\n Some(value) => {\n output.insert(key, value.clone());\n }\n None => {\n if let Some(value) = agent_variable.default.clone() {\n output.insert(key, value);\n continue;\n }\n if no_interaction {\n continue;\n }\n if *IS_STDOUT_TERMINAL {\n if !printed {\n println!(\"⚙ Init agent variables...\");\n printed = true;\n }\n let value = Text::new(&format!(\n \"{} ({}):\",\n agent_variable.name, agent_variable.description\n ))\n .with_validator(|input: &str| {\n if input.trim().is_empty() {\n Ok(Validation::Invalid(\"This field is required\".into()))\n } else {\n Ok(Validation::Valid)\n }\n })\n .prompt()?;\n output.insert(key, value);\n } else {\n unset_variables.push(agent_variable)\n }\n }\n }\n }\n if !unset_variables.is_empty() {\n bail!(\n \"The following agent variables are required:\\n{}\",\n unset_variables\n .iter()\n .map(|v| format!(\" - {}: {}\", v.name, v.description))\n .collect::>()\n .join(\"\\n\")\n )\n }\n Ok(output)\n }\n\n pub fn export(&self) -> Result {\n let mut value = json!({});\n value[\"name\"] = json!(self.name());\n let variables = self.variables();\n if !variables.is_empty() {\n value[\"variables\"] = serde_json::to_value(variables)?;\n }\n value[\"config\"] = json!(self.config);\n let mut definition = self.definition.clone();\n definition.instructions = self.interpolated_instructions();\n value[\"definition\"] = json!(definition);\n value[\"functions_dir\"] = Config::agent_functions_dir(&self.name)\n .display()\n .to_string()\n .into();\n value[\"data_dir\"] = Config::agent_data_dir(&self.name)\n .display()\n .to_string()\n .into();\n value[\"config_file\"] = Config::agent_config_file(&self.name)\n .display()\n .to_string()\n .into();\n let data = serde_yaml::to_string(&value)?;\n Ok(data)\n }\n\n pub fn banner(&self) -> String {\n self.definition.banner()\n }\n\n pub fn name(&self) -> &str {\n &self.name\n }\n\n pub fn functions(&self) -> &Functions {\n &self.functions\n }\n\n pub fn rag(&self) -> Option> {\n self.rag.clone()\n }\n\n pub fn conversation_staters(&self) -> &[String] {\n &self.definition.conversation_starters\n }\n\n pub fn interpolated_instructions(&self) -> String {\n let mut output = self\n .session_dynamic_instructions\n .clone()\n .or_else(|| self.shared_dynamic_instructions.clone())\n .or_else(|| self.config.instructions.clone())\n .unwrap_or_else(|| self.definition.instructions.clone());\n for (k, v) in self.variables() {\n output = output.replace(&format!(\"{{{{{k}}}}}\"), v)\n }\n interpolate_variables(&mut output);\n output\n }\n\n pub fn agent_prelude(&self) -> Option<&str> {\n self.config.agent_prelude.as_deref()\n }\n\n pub fn variables(&self) -> &AgentVariables {\n match &self.session_variables {\n Some(variables) => variables,\n None => &self.shared_variables,\n }\n }\n\n pub fn variable_envs(&self) -> HashMap {\n self.variables()\n .iter()\n .map(|(k, v)| {\n (\n format!(\"LLM_AGENT_VAR_{}\", normalize_env_name(k)),\n v.clone(),\n )\n })\n .collect()\n }\n\n pub fn config_variables(&self) -> &AgentVariables {\n &self.config.variables\n }\n\n pub fn shared_variables(&self) -> &AgentVariables {\n &self.shared_variables\n }\n\n pub fn set_shared_variables(&mut self, shared_variables: AgentVariables) {\n self.shared_variables = shared_variables;\n }\n\n pub fn set_session_variables(&mut self, session_variables: AgentVariables) {\n self.session_variables = Some(session_variables);\n }\n\n pub fn defined_variables(&self) -> &[AgentVariable] {\n &self.definition.variables\n }\n\n pub fn exit_session(&mut self) {\n self.session_variables = None;\n self.session_dynamic_instructions = None;\n }\n\n pub fn is_dynamic_instructions(&self) -> bool {\n self.definition.dynamic_instructions\n }\n\n pub fn update_shared_dynamic_instructions(&mut self, force: bool) -> Result<()> {\n if self.is_dynamic_instructions() && (force || self.shared_dynamic_instructions.is_none()) {\n self.shared_dynamic_instructions = Some(self.run_instructions_fn()?);\n }\n Ok(())\n }\n\n pub fn update_session_dynamic_instructions(&mut self, value: Option) -> Result<()> {\n if self.is_dynamic_instructions() {\n let value = match value {\n Some(v) => v,\n None => self.run_instructions_fn()?,\n };\n self.session_dynamic_instructions = Some(value);\n }\n Ok(())\n }\n\n fn run_instructions_fn(&self) -> Result {\n let value = run_llm_function(\n self.name().to_string(),\n vec![\"_instructions\".into(), \"{}\".into()],\n self.variable_envs(),\n )?;\n match value {\n Some(v) => Ok(v),\n _ => bail!(\"No return value from '_instructions' function\"),\n }\n }\n}\n\nimpl RoleLike for Agent {\n fn to_role(&self) -> Role {\n let prompt = self.interpolated_instructions();\n let mut role = Role::new(\"\", &prompt);\n role.sync(self);\n role\n }\n\n fn model(&self) -> &Model {\n &self.model\n }\n\n fn temperature(&self) -> Option {\n self.config.temperature\n }\n\n fn top_p(&self) -> Option {\n self.config.top_p\n }\n\n fn use_tools(&self) -> Option {\n self.config.use_tools.clone()\n }\n\n fn set_model(&mut self, model: Model) {\n self.config.model_id = Some(model.id());\n self.model = model;\n }\n\n fn set_temperature(&mut self, value: Option) {\n self.config.temperature = value;\n }\n\n fn set_top_p(&mut self, value: Option) {\n self.config.top_p = value;\n }\n\n fn set_use_tools(&mut self, value: Option) {\n self.config.use_tools = value;\n }\n}\n\n#[derive(Debug, Clone, Default, Deserialize, Serialize)]\npub struct AgentConfig {\n #[serde(rename(serialize = \"model\", deserialize = \"model\"))]\n pub model_id: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub temperature: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub top_p: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub use_tools: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub agent_prelude: Option,\n #[serde(default, skip_serializing_if = \"Option::is_none\")]\n pub instructions: Option,\n #[serde(default, skip_serializing_if = \"IndexMap::is_empty\")]\n pub variables: AgentVariables,\n}\n\nimpl AgentConfig {\n pub fn new(config: &Config) -> Self {\n Self {\n use_tools: config.use_tools.clone(),\n agent_prelude: config.agent_prelude.clone(),\n ..Default::default()\n }\n }\n\n pub fn load(path: &Path) -> Result {\n let contents = read_to_string(path)\n .with_context(|| format!(\"Failed to read agent config file at '{}'\", path.display()))?;\n let config: Self = serde_yaml::from_str(&contents)\n .with_context(|| format!(\"Failed to load agent config at '{}'\", path.display()))?;\n Ok(config)\n }\n\n fn load_envs(&mut self, name: &str) {\n let with_prefix = |v: &str| normalize_env_name(&format!(\"{name}_{v}\"));\n\n if let Some(v) = read_env_value::(&with_prefix(\"model\")) {\n self.model_id = v;\n }\n if let Some(v) = read_env_value::(&with_prefix(\"temperature\")) {\n self.temperature = v;\n }\n if let Some(v) = read_env_value::(&with_prefix(\"top_p\")) {\n self.top_p = v;\n }\n if let Some(v) = read_env_value::(&with_prefix(\"use_tools\")) {\n self.use_tools = v;\n }\n if let Some(v) = read_env_value::(&with_prefix(\"agent_prelude\")) {\n self.agent_prelude = v;\n }\n if let Some(v) = read_env_value::(&with_prefix(\"instructions\")) {\n self.instructions = v;\n }\n if let Ok(v) = env::var(with_prefix(\"variables\")) {\n if let Ok(v) = serde_json::from_str(&v) {\n self.variables = v;\n }\n }\n }\n}\n\n#[derive(Debug, Clone, Default, Deserialize, Serialize)]\npub struct AgentDefinition {\n pub name: String,\n #[serde(default)]\n pub description: String,\n #[serde(default)]\n pub version: String,\n #[serde(default)]\n pub instructions: String,\n #[serde(default)]\n pub dynamic_instructions: bool,\n #[serde(default)]\n pub variables: Vec,\n #[serde(default)]\n pub conversation_starters: Vec,\n #[serde(default)]\n pub documents: Vec,\n}\n\nimpl AgentDefinition {\n pub fn load(path: &Path) -> Result {\n let contents = read_to_string(path)\n .with_context(|| format!(\"Failed to read agent index file at '{}'\", path.display()))?;\n let definition: Self = serde_yaml::from_str(&contents)\n .with_context(|| format!(\"Failed to load agent index at '{}'\", path.display()))?;\n Ok(definition)\n }\n\n fn banner(&self) -> String {\n let AgentDefinition {\n name,\n description,\n version,\n conversation_starters,\n ..\n } = self;\n let starters = if conversation_starters.is_empty() {\n String::new()\n } else {\n let starters = conversation_starters\n .iter()\n .map(|v| format!(\"- {v}\"))\n .collect::>()\n .join(\"\\n\");\n format!(\n r#\"\n\n## Conversation Starters\n{starters}\"#\n )\n };\n format!(\n r#\"# {name} {version}\n{description}{starters}\"#\n )\n }\n\n fn replace_tools_placeholder(&mut self, functions: &Functions) {\n let tools_placeholder: &str = \"{{__tools__}}\";\n if self.instructions.contains(tools_placeholder) {\n let tools = functions\n .declarations()\n .iter()\n .enumerate()\n .map(|(i, v)| {\n let description = match v.description.split_once('\\n') {\n Some((v, _)) => v,\n None => &v.description,\n };\n format!(\"{}. {}: {description}\", i + 1, v.name)\n })\n .collect::>()\n .join(\"\\n\");\n self.instructions = self.instructions.replace(tools_placeholder, &tools);\n }\n }\n}\n\n#[derive(Debug, Clone, Default, Deserialize, Serialize)]\npub struct AgentVariable {\n pub name: String,\n pub description: String,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub default: Option,\n #[serde(skip_deserializing, default)]\n pub value: String,\n}\n\npub fn list_agents() -> Vec {\n let agents_file = Config::functions_dir().join(\"agents.txt\");\n let contents = match read_to_string(agents_file) {\n Ok(v) => v,\n Err(_) => return vec![],\n };\n contents\n .split('\\n')\n .filter_map(|line| {\n let line = line.trim();\n if line.is_empty() || line.starts_with('#') {\n None\n } else {\n Some(line.to_string())\n }\n })\n .collect()\n}\n\npub fn complete_agent_variables(agent_name: &str) -> Vec<(String, Option)> {\n let index_path = Config::agent_functions_dir(agent_name).join(\"index.yaml\");\n if !index_path.exists() {\n return vec![];\n }\n let Ok(definition) = AgentDefinition::load(&index_path) else {\n return vec![];\n };\n definition\n .variables\n .iter()\n .map(|v| {\n let description = match &v.default {\n Some(default) => format!(\"{} [default: {default}]\", v.description),\n None => v.description.clone(),\n };\n (format!(\"{}=\", v.name), Some(description))\n })\n .collect()\n}\n" }, { "path": "src/config/input.rs", "content": "use super::*;\n\nuse crate::client::{\n init_client, patch_messages, ChatCompletionsData, Client, ImageUrl, Message, MessageContent,\n MessageContentPart, MessageContentToolCalls, MessageRole, Model,\n};\nuse crate::function::ToolResult;\nuse crate::utils::{base64_encode, is_loader_protocol, sha256, AbortSignal};\n\nuse anyhow::{bail, Context, Result};\nuse indexmap::IndexSet;\nuse std::{collections::HashMap, fs::File, io::Read};\nuse unicode_width::{UnicodeWidthChar, UnicodeWidthStr};\n\nconst IMAGE_EXTS: [&str; 5] = [\"png\", \"jpeg\", \"jpg\", \"webp\", \"gif\"];\nconst SUMMARY_MAX_WIDTH: usize = 80;\n\n#[derive(Debug, Clone)]\npub struct Input {\n config: GlobalConfig,\n text: String,\n raw: (String, Vec),\n patched_text: Option,\n last_reply: Option,\n continue_output: Option,\n regenerate: bool,\n medias: Vec,\n data_urls: HashMap,\n tool_calls: Option,\n role: Role,\n rag_name: Option,\n with_session: bool,\n with_agent: bool,\n}\n\nimpl Input {\n pub fn from_str(config: &GlobalConfig, text: &str, role: Option) -> Self {\n let (role, with_session, with_agent) = resolve_role(&config.read(), role);\n Self {\n config: config.clone(),\n text: text.to_string(),\n raw: (text.to_string(), vec![]),\n patched_text: None,\n last_reply: None,\n continue_output: None,\n regenerate: false,\n medias: Default::default(),\n data_urls: Default::default(),\n tool_calls: None,\n role,\n rag_name: None,\n with_session,\n with_agent,\n }\n }\n\n pub async fn from_files(\n config: &GlobalConfig,\n raw_text: &str,\n paths: Vec,\n role: Option,\n ) -> Result {\n let loaders = config.read().document_loaders.clone();\n let (raw_paths, local_paths, remote_urls, external_cmds, protocol_paths, with_last_reply) =\n resolve_paths(&loaders, paths)?;\n let mut last_reply = None;\n let (documents, medias, data_urls) = load_documents(\n &loaders,\n local_paths,\n remote_urls,\n external_cmds,\n protocol_paths,\n )\n .await\n .context(\"Failed to load files\")?;\n let mut texts = vec![];\n if !raw_text.is_empty() {\n texts.push(raw_text.to_string());\n };\n if with_last_reply {\n if let Some(LastMessage { input, output, .. }) = config.read().last_message.as_ref() {\n if !output.is_empty() {\n last_reply = Some(output.clone())\n } else if let Some(v) = input.last_reply.as_ref() {\n last_reply = Some(v.clone());\n }\n if let Some(v) = last_reply.clone() {\n texts.push(format!(\"\\n{v}\"));\n }\n }\n if last_reply.is_none() && documents.is_empty() && medias.is_empty() {\n bail!(\"No last reply found\");\n }\n }\n let documents_len = documents.len();\n for (kind, path, contents) in documents {\n if documents_len == 1 && raw_text.is_empty() {\n texts.push(format!(\"\\n{contents}\"));\n } else {\n texts.push(format!(\n \"\\n============ {kind}: {path} ============\\n{contents}\"\n ));\n }\n }\n let (role, with_session, with_agent) = resolve_role(&config.read(), role);\n Ok(Self {\n config: config.clone(),\n text: texts.join(\"\\n\"),\n raw: (raw_text.to_string(), raw_paths),\n patched_text: None,\n last_reply,\n continue_output: None,\n regenerate: false,\n medias,\n data_urls,\n tool_calls: Default::default(),\n role,\n rag_name: None,\n with_session,\n with_agent,\n })\n }\n\n pub async fn from_files_with_spinner(\n config: &GlobalConfig,\n raw_text: &str,\n paths: Vec,\n role: Option,\n abort_signal: AbortSignal,\n ) -> Result {\n abortable_run_with_spinner(\n Input::from_files(config, raw_text, paths, role),\n \"Loading files\",\n abort_signal,\n )\n .await\n }\n\n pub fn is_empty(&self) -> bool {\n self.text.is_empty() && self.medias.is_empty()\n }\n\n pub fn data_urls(&self) -> HashMap {\n self.data_urls.clone()\n }\n\n pub fn tool_calls(&self) -> &Option {\n &self.tool_calls\n }\n\n pub fn text(&self) -> String {\n match self.patched_text.clone() {\n Some(text) => text,\n None => self.text.clone(),\n }\n }\n\n pub fn clear_patch(&mut self) {\n self.patched_text = None;\n }\n\n pub fn set_text(&mut self, text: String) {\n self.text = text;\n }\n\n pub fn stream(&self) -> bool {\n self.config.read().stream && !self.role().model().no_stream()\n }\n\n pub fn continue_output(&self) -> Option<&str> {\n self.continue_output.as_deref()\n }\n\n pub fn set_continue_output(&mut self, output: &str) {\n let output = match &self.continue_output {\n Some(v) => format!(\"{v}{output}\"),\n None => output.to_string(),\n };\n self.continue_output = Some(output);\n }\n\n pub fn regenerate(&self) -> bool {\n self.regenerate\n }\n\n pub fn set_regenerate(&mut self) {\n let role = self.config.read().extract_role();\n if role.name() == self.role().name() {\n self.role = role;\n }\n self.regenerate = true;\n self.tool_calls = None;\n }\n\n pub async fn use_embeddings(&mut self, abort_signal: AbortSignal) -> Result<()> {\n if self.text.is_empty() {\n return Ok(());\n }\n let rag = self.config.read().rag.clone();\n if let Some(rag) = rag {\n let result = Config::search_rag(&self.config, &rag, &self.text, abort_signal).await?;\n self.patched_text = Some(result);\n self.rag_name = Some(rag.name().to_string());\n }\n Ok(())\n }\n\n pub fn rag_name(&self) -> Option<&str> {\n self.rag_name.as_deref()\n }\n\n pub fn merge_tool_results(mut self, output: String, tool_results: Vec) -> Self {\n match self.tool_calls.as_mut() {\n Some(exist_tool_results) => {\n exist_tool_results.merge(tool_results, output);\n }\n None => self.tool_calls = Some(MessageContentToolCalls::new(tool_results, output)),\n }\n self\n }\n\n pub fn create_client(&self) -> Result> {\n init_client(&self.config, Some(self.role().model().clone()))\n }\n\n pub async fn fetch_chat_text(&self) -> Result {\n let client = self.create_client()?;\n let text = client.chat_completions(self.clone()).await?.text;\n let text = strip_think_tag(&text).to_string();\n Ok(text)\n }\n\n pub fn prepare_completion_data(\n &self,\n model: &Model,\n stream: bool,\n ) -> Result {\n let mut messages = self.build_messages()?;\n patch_messages(&mut messages, model);\n model.guard_max_input_tokens(&messages)?;\n let (temperature, top_p) = (self.role().temperature(), self.role().top_p());\n let functions = self.config.read().select_functions(self.role());\n Ok(ChatCompletionsData {\n messages,\n temperature,\n top_p,\n functions,\n stream,\n })\n }\n\n pub fn build_messages(&self) -> Result> {\n let mut messages = if let Some(session) = self.session(&self.config.read().session) {\n session.build_messages(self)\n } else {\n self.role().build_messages(self)\n };\n if let Some(tool_calls) = &self.tool_calls {\n messages.push(Message::new(\n MessageRole::Assistant,\n MessageContent::ToolCalls(tool_calls.clone()),\n ))\n }\n Ok(messages)\n }\n\n pub fn echo_messages(&self) -> String {\n if let Some(session) = self.session(&self.config.read().session) {\n session.echo_messages(self)\n } else {\n self.role().echo_messages(self)\n }\n }\n\n pub fn role(&self) -> &Role {\n &self.role\n }\n\n pub fn session<'a>(&self, session: &'a Option) -> Option<&'a Session> {\n if self.with_session {\n session.as_ref()\n } else {\n None\n }\n }\n\n pub fn session_mut<'a>(&self, session: &'a mut Option) -> Option<&'a mut Session> {\n if self.with_session {\n session.as_mut()\n } else {\n None\n }\n }\n\n pub fn with_agent(&self) -> bool {\n self.with_agent\n }\n\n pub fn summary(&self) -> String {\n let text: String = self\n .text\n .trim()\n .chars()\n .map(|c| if c.is_control() { ' ' } else { c })\n .collect();\n if text.width_cjk() > SUMMARY_MAX_WIDTH {\n let mut sum_width = 0;\n let mut chars = vec![];\n for c in text.chars() {\n sum_width += c.width_cjk().unwrap_or(1);\n if sum_width > SUMMARY_MAX_WIDTH - 3 {\n chars.extend(['.', '.', '.']);\n break;\n }\n chars.push(c);\n }\n chars.into_iter().collect()\n } else {\n text\n }\n }\n\n pub fn raw(&self) -> String {\n let (text, files) = &self.raw;\n let mut segments = files.to_vec();\n if !segments.is_empty() {\n segments.insert(0, \".file\".into());\n }\n if !text.is_empty() {\n if !segments.is_empty() {\n segments.push(\"--\".into());\n }\n segments.push(text.clone());\n }\n segments.join(\" \")\n }\n\n pub fn render(&self) -> String {\n let text = self.text();\n if self.medias.is_empty() {\n return text;\n }\n let tail_text = if text.is_empty() {\n String::new()\n } else {\n format!(\" -- {text}\")\n };\n let files: Vec = self\n .medias\n .iter()\n .cloned()\n .map(|url| resolve_data_url(&self.data_urls, url))\n .collect();\n format!(\".file {}{}\", files.join(\" \"), tail_text)\n }\n\n pub fn message_content(&self) -> MessageContent {\n if self.medias.is_empty() {\n MessageContent::Text(self.text())\n } else {\n let mut list: Vec = self\n .medias\n .iter()\n .cloned()\n .map(|url| MessageContentPart::ImageUrl {\n image_url: ImageUrl { url },\n })\n .collect();\n if !self.text.is_empty() {\n list.insert(0, MessageContentPart::Text { text: self.text() });\n }\n MessageContent::Array(list)\n }\n }\n}\n\nfn resolve_role(config: &Config, role: Option) -> (Role, bool, bool) {\n match role {\n Some(v) => (v, false, false),\n None => (\n config.extract_role(),\n config.session.is_some(),\n config.agent.is_some(),\n ),\n }\n}\n\ntype ResolvePathsOutput = (\n Vec,\n Vec,\n Vec,\n Vec,\n Vec,\n bool,\n);\n\nfn resolve_paths(\n loaders: &HashMap,\n paths: Vec,\n) -> Result {\n let mut raw_paths = IndexSet::new();\n let mut local_paths = IndexSet::new();\n let mut remote_urls = IndexSet::new();\n let mut external_cmds = IndexSet::new();\n let mut protocol_paths = IndexSet::new();\n let mut with_last_reply = false;\n for path in paths {\n if path == \"%%\" {\n with_last_reply = true;\n raw_paths.insert(path);\n } else if path.starts_with('`') && path.len() > 2 && path.ends_with('`') {\n external_cmds.insert(path[1..path.len() - 1].to_string());\n raw_paths.insert(path);\n } else if is_url(&path) {\n if path.strip_suffix(\"**\").is_some() {\n bail!(\"Invalid website '{path}'\");\n }\n remote_urls.insert(path.clone());\n raw_paths.insert(path);\n } else if is_loader_protocol(loaders, &path) {\n protocol_paths.insert(path.clone());\n raw_paths.insert(path);\n } else {\n let resolved_path = resolve_home_dir(&path);\n let absolute_path = to_absolute_path(&resolved_path)\n .with_context(|| format!(\"Invalid path '{path}'\"))?;\n local_paths.insert(resolved_path);\n raw_paths.insert(absolute_path);\n }\n }\n Ok((\n raw_paths.into_iter().collect(),\n local_paths.into_iter().collect(),\n remote_urls.into_iter().collect(),\n external_cmds.into_iter().collect(),\n protocol_paths.into_iter().collect(),\n with_last_reply,\n ))\n}\n\nasync fn load_documents(\n loaders: &HashMap,\n local_paths: Vec,\n remote_urls: Vec,\n external_cmds: Vec,\n protocol_paths: Vec,\n) -> Result<(\n Vec<(&'static str, String, String)>,\n Vec,\n HashMap,\n)> {\n let mut files = vec![];\n let mut medias = vec![];\n let mut data_urls = HashMap::new();\n\n for cmd in external_cmds {\n let output = duct::cmd(&SHELL.cmd, &[&SHELL.arg, &cmd])\n .stderr_to_stdout()\n .unchecked()\n .read()\n .unwrap_or_else(|err| err.to_string());\n files.push((\"CMD\", cmd, output));\n }\n\n let local_files = expand_glob_paths(&local_paths, true).await?;\n for file_path in local_files {\n if is_image(&file_path) {\n let contents = read_media_to_data_url(&file_path)\n .with_context(|| format!(\"Unable to read media '{file_path}'\"))?;\n data_urls.insert(sha256(&contents), file_path);\n medias.push(contents)\n } else {\n let document = load_file(loaders, &file_path)\n .await\n .with_context(|| format!(\"Unable to read file '{file_path}'\"))?;\n files.push((\"FILE\", file_path, document.contents));\n }\n }\n\n for file_url in remote_urls {\n let (contents, extension) = fetch_with_loaders(loaders, &file_url, true)\n .await\n .with_context(|| format!(\"Failed to load url '{file_url}'\"))?;\n if extension == MEDIA_URL_EXTENSION {\n data_urls.insert(sha256(&contents), file_url);\n medias.push(contents)\n } else {\n files.push((\"URL\", file_url, contents));\n }\n }\n\n for protocol_path in protocol_paths {\n let documents = load_protocol_path(loaders, &protocol_path)\n .with_context(|| format!(\"Failed to load from '{protocol_path}'\"))?;\n files.extend(\n documents\n .into_iter()\n .map(|document| (\"FROM\", document.path, document.contents)),\n );\n }\n\n Ok((files, medias, data_urls))\n}\n\npub fn resolve_data_url(data_urls: &HashMap, data_url: String) -> String {\n if data_url.starts_with(\"data:\") {\n let hash = sha256(&data_url);\n if let Some(path) = data_urls.get(&hash) {\n return path.to_string();\n }\n data_url\n } else {\n data_url\n }\n}\n\nfn is_image(path: &str) -> bool {\n get_patch_extension(path)\n .map(|v| IMAGE_EXTS.contains(&v.as_str()))\n .unwrap_or_default()\n}\n\nfn read_media_to_data_url(image_path: &str) -> Result {\n let extension = get_patch_extension(image_path).unwrap_or_default();\n let mime_type = match extension.as_str() {\n \"png\" => \"image/png\",\n \"jpg\" | \"jpeg\" => \"image/jpeg\",\n \"webp\" => \"image/webp\",\n \"gif\" => \"image/gif\",\n _ => bail!(\"Unexpected media type\"),\n };\n let mut file = File::open(image_path)?;\n let mut buffer = Vec::new();\n file.read_to_end(&mut buffer)?;\n\n let encoded_image = base64_encode(buffer);\n let data_url = format!(\"data:{mime_type};base64,{encoded_image}\");\n\n Ok(data_url)\n}\n" }, { "path": "src/config/mod.rs", "content": "mod agent;\nmod input;\nmod role;\nmod session;\n\npub use self::agent::{complete_agent_variables, list_agents, Agent, AgentVariables};\npub use self::input::Input;\npub use self::role::{\n Role, RoleLike, CODE_ROLE, CREATE_TITLE_ROLE, EXPLAIN_SHELL_ROLE, SHELL_ROLE,\n};\nuse self::session::Session;\n\nuse crate::client::{\n create_client_config, list_client_types, list_models, ClientConfig, MessageContentToolCalls,\n Model, ModelType, ProviderModels, OPENAI_COMPATIBLE_PROVIDERS,\n};\nuse crate::function::{FunctionDeclaration, Functions, ToolResult};\nuse crate::rag::Rag;\nuse crate::render::{MarkdownRender, RenderOptions};\nuse crate::repl::{run_repl_command, split_args_text};\nuse crate::utils::*;\n\nuse anyhow::{anyhow, bail, Context, Result};\nuse indexmap::IndexMap;\nuse inquire::{list_option::ListOption, validator::Validation, Confirm, MultiSelect, Select, Text};\nuse parking_lot::RwLock;\nuse serde::{Deserialize, Serialize};\nuse serde_json::json;\nuse simplelog::LevelFilter;\nuse std::collections::{HashMap, HashSet};\nuse std::{\n env,\n fs::{\n create_dir_all, read_dir, read_to_string, remove_dir_all, remove_file, File, OpenOptions,\n },\n io::Write,\n path::{Path, PathBuf},\n process,\n sync::{Arc, OnceLock},\n};\nuse syntect::highlighting::ThemeSet;\nuse terminal_colorsaurus::{color_scheme, ColorScheme, QueryOptions};\n\npub const TEMP_ROLE_NAME: &str = \"%%\";\npub const TEMP_RAG_NAME: &str = \"temp\";\npub const TEMP_SESSION_NAME: &str = \"temp\";\n\n/// Monokai Extended\nconst DARK_THEME: &[u8] = include_bytes!(\"../../assets/monokai-extended.theme.bin\");\nconst LIGHT_THEME: &[u8] = include_bytes!(\"../../assets/monokai-extended-light.theme.bin\");\n\nconst CONFIG_FILE_NAME: &str = \"config.yaml\";\nconst ROLES_DIR_NAME: &str = \"roles\";\nconst MACROS_DIR_NAME: &str = \"macros\";\nconst ENV_FILE_NAME: &str = \".env\";\nconst MESSAGES_FILE_NAME: &str = \"messages.md\";\nconst SESSIONS_DIR_NAME: &str = \"sessions\";\nconst RAGS_DIR_NAME: &str = \"rags\";\nconst FUNCTIONS_DIR_NAME: &str = \"functions\";\nconst FUNCTIONS_FILE_NAME: &str = \"functions.json\";\nconst FUNCTIONS_BIN_DIR_NAME: &str = \"bin\";\nconst AGENTS_DIR_NAME: &str = \"agents\";\n\nconst CLIENTS_FIELD: &str = \"clients\";\n\nconst SERVE_ADDR: &str = \"127.0.0.1:8000\";\n\nconst SYNC_MODELS_URL: &str =\n \"https://raw.githubusercontent.com/sigoden/aichat/refs/heads/main/models.yaml\";\n\nconst SUMMARIZE_PROMPT: &str =\n \"Summarize the discussion briefly in 200 words or less to use as a prompt for future context.\";\nconst SUMMARY_PROMPT: &str = \"This is a summary of the chat history as a recap: \";\n\nconst RAG_TEMPLATE: &str = r#\"Answer the query based on the context while respecting the rules. (user query, some textual context and rules, all inside xml tags)\n\n\n__CONTEXT__\n\n\n\n- If you don't know, just say so.\n- If you are not sure, ask for clarification.\n- Answer in the same language as the user query.\n- If the context appears unreadable or of poor quality, tell the user then answer as best as you can.\n- If the answer is not in the context but you think you know the answer, explain that to the user then answer with your own knowledge.\n- Answer directly and without using xml tags.\n\n\n\n__INPUT__\n\"#;\n\nconst LEFT_PROMPT: &str = \"{color.green}{?session {?agent {agent}>}{session}{?role /}}{!session {?agent {agent}>}}{role}{?rag @{rag}}{color.cyan}{?session )}{!session >}{color.reset} \";\nconst RIGHT_PROMPT: &str = \"{color.purple}{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}\";\n\nstatic EDITOR: OnceLock> = OnceLock::new();\n\n#[derive(Debug, Clone, Deserialize)]\n#[serde(default)]\npub struct Config {\n #[serde(rename(serialize = \"model\", deserialize = \"model\"))]\n #[serde(default)]\n pub model_id: String,\n pub temperature: Option,\n pub top_p: Option,\n\n pub dry_run: bool,\n pub stream: bool,\n pub save: bool,\n pub keybindings: String,\n pub editor: Option,\n pub wrap: Option,\n pub wrap_code: bool,\n\n pub function_calling: bool,\n pub mapping_tools: IndexMap,\n pub use_tools: Option,\n\n pub repl_prelude: Option,\n pub cmd_prelude: Option,\n pub agent_prelude: Option,\n\n pub save_session: Option,\n pub compress_threshold: usize,\n pub summarize_prompt: Option,\n pub summary_prompt: Option,\n\n pub rag_embedding_model: Option,\n pub rag_reranker_model: Option,\n pub rag_top_k: usize,\n pub rag_chunk_size: Option,\n pub rag_chunk_overlap: Option,\n pub rag_template: Option,\n\n #[serde(default)]\n pub document_loaders: HashMap,\n\n pub highlight: bool,\n pub theme: Option,\n pub left_prompt: Option,\n pub right_prompt: Option,\n\n pub serve_addr: Option,\n pub user_agent: Option,\n pub save_shell_history: bool,\n pub sync_models_url: Option,\n\n pub clients: Vec,\n\n #[serde(skip)]\n pub macro_flag: bool,\n #[serde(skip)]\n pub info_flag: bool,\n #[serde(skip)]\n pub agent_variables: Option,\n\n #[serde(skip)]\n pub model: Model,\n #[serde(skip)]\n pub functions: Functions,\n #[serde(skip)]\n pub working_mode: WorkingMode,\n #[serde(skip)]\n pub last_message: Option,\n\n #[serde(skip)]\n pub role: Option,\n #[serde(skip)]\n pub session: Option,\n #[serde(skip)]\n pub rag: Option>,\n #[serde(skip)]\n pub agent: Option,\n}\n\nimpl Default for Config {\n fn default() -> Self {\n Self {\n model_id: Default::default(),\n temperature: None,\n top_p: None,\n\n dry_run: false,\n stream: true,\n save: false,\n keybindings: \"emacs\".into(),\n editor: None,\n wrap: None,\n wrap_code: false,\n\n function_calling: true,\n mapping_tools: Default::default(),\n use_tools: None,\n\n repl_prelude: None,\n cmd_prelude: None,\n agent_prelude: None,\n\n save_session: None,\n compress_threshold: 4000,\n summarize_prompt: None,\n summary_prompt: None,\n\n rag_embedding_model: None,\n rag_reranker_model: None,\n rag_top_k: 5,\n rag_chunk_size: None,\n rag_chunk_overlap: None,\n rag_template: None,\n\n document_loaders: Default::default(),\n\n highlight: true,\n theme: None,\n left_prompt: None,\n right_prompt: None,\n\n serve_addr: None,\n user_agent: None,\n save_shell_history: true,\n sync_models_url: None,\n\n clients: vec![],\n\n macro_flag: false,\n info_flag: false,\n agent_variables: None,\n\n model: Default::default(),\n functions: Default::default(),\n working_mode: WorkingMode::Cmd,\n last_message: None,\n\n role: None,\n session: None,\n rag: None,\n agent: None,\n }\n }\n}\n\npub type GlobalConfig = Arc>;\n\nimpl Config {\n pub async fn init(working_mode: WorkingMode, info_flag: bool) -> Result {\n let config_path = Self::config_file();\n let mut config = if !config_path.exists() {\n match env::var(get_env_name(\"provider\"))\n .ok()\n .or_else(|| env::var(get_env_name(\"platform\")).ok())\n {\n Some(v) => Self::load_dynamic(&v)?,\n None => {\n if *IS_STDOUT_TERMINAL {\n create_config_file(&config_path).await?;\n }\n Self::load_from_file(&config_path)?\n }\n }\n } else {\n Self::load_from_file(&config_path)?\n };\n\n config.working_mode = working_mode;\n config.info_flag = info_flag;\n\n let setup = |config: &mut Self| -> Result<()> {\n config.load_envs();\n\n if let Some(wrap) = config.wrap.clone() {\n config.set_wrap(&wrap)?;\n }\n\n config.load_functions()?;\n\n config.setup_model()?;\n config.setup_document_loaders();\n config.setup_user_agent();\n Ok(())\n };\n let ret = setup(&mut config);\n if !info_flag {\n ret?;\n }\n Ok(config)\n }\n\n pub fn config_dir() -> PathBuf {\n if let Ok(v) = env::var(get_env_name(\"config_dir\")) {\n PathBuf::from(v)\n } else if let Ok(v) = env::var(\"XDG_CONFIG_HOME\") {\n PathBuf::from(v).join(env!(\"CARGO_CRATE_NAME\"))\n } else {\n let dir = dirs::config_dir().expect(\"No user's config directory\");\n dir.join(env!(\"CARGO_CRATE_NAME\"))\n }\n }\n\n pub fn local_path(name: &str) -> PathBuf {\n Self::config_dir().join(name)\n }\n\n pub fn config_file() -> PathBuf {\n match env::var(get_env_name(\"config_file\")) {\n Ok(value) => PathBuf::from(value),\n Err(_) => Self::local_path(CONFIG_FILE_NAME),\n }\n }\n\n pub fn roles_dir() -> PathBuf {\n match env::var(get_env_name(\"roles_dir\")) {\n Ok(value) => PathBuf::from(value),\n Err(_) => Self::local_path(ROLES_DIR_NAME),\n }\n }\n\n pub fn role_file(name: &str) -> PathBuf {\n Self::roles_dir().join(format!(\"{name}.md\"))\n }\n\n pub fn macros_dir() -> PathBuf {\n match env::var(get_env_name(\"macros_dir\")) {\n Ok(value) => PathBuf::from(value),\n Err(_) => Self::local_path(MACROS_DIR_NAME),\n }\n }\n\n pub fn macro_file(name: &str) -> PathBuf {\n Self::macros_dir().join(format!(\"{name}.yaml\"))\n }\n\n pub fn env_file() -> PathBuf {\n match env::var(get_env_name(\"env_file\")) {\n Ok(value) => PathBuf::from(value),\n Err(_) => Self::local_path(ENV_FILE_NAME),\n }\n }\n\n pub fn messages_file(&self) -> PathBuf {\n match &self.agent {\n None => match env::var(get_env_name(\"messages_file\")) {\n Ok(value) => PathBuf::from(value),\n Err(_) => Self::local_path(MESSAGES_FILE_NAME),\n },\n Some(agent) => Self::agent_data_dir(agent.name()).join(MESSAGES_FILE_NAME),\n }\n }\n\n pub fn sessions_dir(&self) -> PathBuf {\n match &self.agent {\n None => match env::var(get_env_name(\"sessions_dir\")) {\n Ok(value) => PathBuf::from(value),\n Err(_) => Self::local_path(SESSIONS_DIR_NAME),\n },\n Some(agent) => Self::agent_data_dir(agent.name()).join(SESSIONS_DIR_NAME),\n }\n }\n\n pub fn rags_dir() -> PathBuf {\n match env::var(get_env_name(\"rags_dir\")) {\n Ok(value) => PathBuf::from(value),\n Err(_) => Self::local_path(RAGS_DIR_NAME),\n }\n }\n\n pub fn functions_dir() -> PathBuf {\n match env::var(get_env_name(\"functions_dir\")) {\n Ok(value) => PathBuf::from(value),\n Err(_) => Self::local_path(FUNCTIONS_DIR_NAME),\n }\n }\n\n pub fn functions_file() -> PathBuf {\n Self::functions_dir().join(FUNCTIONS_FILE_NAME)\n }\n\n pub fn functions_bin_dir() -> PathBuf {\n Self::functions_dir().join(FUNCTIONS_BIN_DIR_NAME)\n }\n\n pub fn session_file(&self, name: &str) -> PathBuf {\n match name.split_once(\"/\") {\n Some((dir, name)) => self.sessions_dir().join(dir).join(format!(\"{name}.yaml\")),\n None => self.sessions_dir().join(format!(\"{name}.yaml\")),\n }\n }\n\n pub fn rag_file(&self, name: &str) -> PathBuf {\n match &self.agent {\n Some(agent) => Self::agent_rag_file(agent.name(), name),\n None => Self::rags_dir().join(format!(\"{name}.yaml\")),\n }\n }\n\n pub fn agents_data_dir() -> PathBuf {\n Self::local_path(AGENTS_DIR_NAME)\n }\n\n pub fn agent_data_dir(name: &str) -> PathBuf {\n match env::var(format!(\"{}_DATA_DIR\", normalize_env_name(name))) {\n Ok(value) => PathBuf::from(value),\n Err(_) => Self::agents_data_dir().join(name),\n }\n }\n\n pub fn agent_config_file(name: &str) -> PathBuf {\n match env::var(format!(\"{}_CONFIG_FILE\", normalize_env_name(name))) {\n Ok(value) => PathBuf::from(value),\n Err(_) => Self::agent_data_dir(name).join(CONFIG_FILE_NAME),\n }\n }\n\n pub fn agent_rag_file(agent_name: &str, rag_name: &str) -> PathBuf {\n Self::agent_data_dir(agent_name).join(format!(\"{rag_name}.yaml\"))\n }\n\n pub fn agents_functions_dir() -> PathBuf {\n Self::functions_dir().join(AGENTS_DIR_NAME)\n }\n\n pub fn agent_functions_dir(name: &str) -> PathBuf {\n match env::var(format!(\"{}_FUNCTIONS_DIR\", normalize_env_name(name))) {\n Ok(value) => PathBuf::from(value),\n Err(_) => Self::agents_functions_dir().join(name),\n }\n }\n\n pub fn models_override_file() -> PathBuf {\n Self::local_path(\"models-override.yaml\")\n }\n\n pub fn state(&self) -> StateFlags {\n let mut flags = StateFlags::empty();\n if let Some(session) = &self.session {\n if session.is_empty() {\n flags |= StateFlags::SESSION_EMPTY;\n } else {\n flags |= StateFlags::SESSION;\n }\n if session.role_name().is_some() {\n flags |= StateFlags::ROLE;\n }\n } else if self.role.is_some() {\n flags |= StateFlags::ROLE;\n }\n if self.agent.is_some() {\n flags |= StateFlags::AGENT;\n }\n if self.rag.is_some() {\n flags |= StateFlags::RAG;\n }\n flags\n }\n\n pub fn serve_addr(&self) -> String {\n self.serve_addr.clone().unwrap_or_else(|| SERVE_ADDR.into())\n }\n\n pub fn log_config(is_serve: bool) -> Result<(LevelFilter, Option)> {\n let log_level = env::var(get_env_name(\"log_level\"))\n .ok()\n .and_then(|v| v.parse().ok())\n .unwrap_or(match cfg!(debug_assertions) {\n true => LevelFilter::Debug,\n false => {\n if is_serve {\n LevelFilter::Info\n } else {\n LevelFilter::Off\n }\n }\n });\n if log_level == LevelFilter::Off {\n return Ok((log_level, None));\n }\n let log_path = match env::var(get_env_name(\"log_path\")) {\n Ok(v) => Some(PathBuf::from(v)),\n Err(_) => match is_serve {\n true => None,\n false => Some(Config::local_path(&format!(\n \"{}.log\",\n env!(\"CARGO_CRATE_NAME\")\n ))),\n },\n };\n Ok((log_level, log_path))\n }\n\n pub fn edit_config(&self) -> Result<()> {\n let config_path = Self::config_file();\n let editor = self.editor()?;\n edit_file(&editor, &config_path)?;\n println!(\n \"NOTE: Remember to restart {} if there are changes made to '{}\",\n env!(\"CARGO_CRATE_NAME\"),\n config_path.display(),\n );\n Ok(())\n }\n\n pub fn current_model(&self) -> &Model {\n if let Some(session) = self.session.as_ref() {\n session.model()\n } else if let Some(agent) = self.agent.as_ref() {\n agent.model()\n } else if let Some(role) = self.role.as_ref() {\n role.model()\n } else {\n &self.model\n }\n }\n\n pub fn role_like_mut(&mut self) -> Option<&mut dyn RoleLike> {\n if let Some(session) = self.session.as_mut() {\n Some(session)\n } else if let Some(agent) = self.agent.as_mut() {\n Some(agent)\n } else if let Some(role) = self.role.as_mut() {\n Some(role)\n } else {\n None\n }\n }\n\n pub fn extract_role(&self) -> Role {\n if let Some(session) = self.session.as_ref() {\n session.to_role()\n } else if let Some(agent) = self.agent.as_ref() {\n agent.to_role()\n } else if let Some(role) = self.role.as_ref() {\n role.clone()\n } else {\n let mut role = Role::default();\n role.batch_set(\n &self.model,\n self.temperature,\n self.top_p,\n self.use_tools.clone(),\n );\n role\n }\n }\n\n pub fn info(&self) -> Result {\n if let Some(agent) = &self.agent {\n let output = agent.export()?;\n if let Some(session) = &self.session {\n let session = session\n .export()?\n .split('\\n')\n .map(|v| format!(\" {v}\"))\n .collect::>()\n .join(\"\\n\");\n Ok(format!(\"{output}session:\\n{session}\"))\n } else {\n Ok(output)\n }\n } else if let Some(session) = &self.session {\n session.export()\n } else if let Some(role) = &self.role {\n Ok(role.export())\n } else if let Some(rag) = &self.rag {\n rag.export()\n } else {\n self.sysinfo()\n }\n }\n\n pub fn sysinfo(&self) -> Result {\n let display_path = |path: &Path| path.display().to_string();\n let wrap = self\n .wrap\n .clone()\n .map_or_else(|| String::from(\"no\"), |v| v.to_string());\n let (rag_reranker_model, rag_top_k) = match &self.rag {\n Some(rag) => rag.get_config(),\n None => (self.rag_reranker_model.clone(), self.rag_top_k),\n };\n let role = self.extract_role();\n let mut items = vec![\n (\"model\", role.model().id()),\n (\"temperature\", format_option_value(&role.temperature())),\n (\"top_p\", format_option_value(&role.top_p())),\n (\"use_tools\", format_option_value(&role.use_tools())),\n (\n \"max_output_tokens\",\n role.model()\n .max_tokens_param()\n .map(|v| format!(\"{v} (current model)\"))\n .unwrap_or_else(|| \"null\".into()),\n ),\n (\"save_session\", format_option_value(&self.save_session)),\n (\"compress_threshold\", self.compress_threshold.to_string()),\n (\n \"rag_reranker_model\",\n format_option_value(&rag_reranker_model),\n ),\n (\"rag_top_k\", rag_top_k.to_string()),\n (\"dry_run\", self.dry_run.to_string()),\n (\"function_calling\", self.function_calling.to_string()),\n (\"stream\", self.stream.to_string()),\n (\"save\", self.save.to_string()),\n (\"keybindings\", self.keybindings.clone()),\n (\"wrap\", wrap),\n (\"wrap_code\", self.wrap_code.to_string()),\n (\"highlight\", self.highlight.to_string()),\n (\"theme\", format_option_value(&self.theme)),\n (\"config_file\", display_path(&Self::config_file())),\n (\"env_file\", display_path(&Self::env_file())),\n (\"roles_dir\", display_path(&Self::roles_dir())),\n (\"sessions_dir\", display_path(&self.sessions_dir())),\n (\"rags_dir\", display_path(&Self::rags_dir())),\n (\"macros_dir\", display_path(&Self::macros_dir())),\n (\"functions_dir\", display_path(&Self::functions_dir())),\n (\"messages_file\", display_path(&self.messages_file())),\n ];\n if let Ok((_, Some(log_path))) = Self::log_config(self.working_mode.is_serve()) {\n items.push((\"log_path\", display_path(&log_path)));\n }\n let output = items\n .iter()\n .map(|(name, value)| format!(\"{name:<24}{value}\\n\"))\n .collect::>()\n .join(\"\");\n Ok(output)\n }\n\n pub fn update(config: &GlobalConfig, data: &str) -> Result<()> {\n let parts: Vec<&str> = data.split_whitespace().collect();\n if parts.len() != 2 {\n bail!(\"Usage: .set . If value is null, unset key.\");\n }\n let key = parts[0];\n let value = parts[1];\n match key {\n \"temperature\" => {\n let value = parse_value(value)?;\n config.write().set_temperature(value);\n }\n \"top_p\" => {\n let value = parse_value(value)?;\n config.write().set_top_p(value);\n }\n \"use_tools\" => {\n let value = parse_value(value)?;\n config.write().set_use_tools(value);\n }\n \"max_output_tokens\" => {\n let value = parse_value(value)?;\n config.write().set_max_output_tokens(value);\n }\n \"save_session\" => {\n let value = parse_value(value)?;\n config.write().set_save_session(value);\n }\n \"compress_threshold\" => {\n let value = parse_value(value)?;\n config.write().set_compress_threshold(value);\n }\n \"rag_reranker_model\" => {\n let value = parse_value(value)?;\n Self::set_rag_reranker_model(config, value)?;\n }\n \"rag_top_k\" => {\n let value = value.parse().with_context(|| \"Invalid value\")?;\n Self::set_rag_top_k(config, value)?;\n }\n \"dry_run\" => {\n let value = value.parse().with_context(|| \"Invalid value\")?;\n config.write().dry_run = value;\n }\n \"function_calling\" => {\n let value = value.parse().with_context(|| \"Invalid value\")?;\n if value && config.write().functions.is_empty() {\n bail!(\"Function calling cannot be enabled because no functions are installed.\")\n }\n config.write().function_calling = value;\n }\n \"stream\" => {\n let value = value.parse().with_context(|| \"Invalid value\")?;\n config.write().stream = value;\n }\n \"save\" => {\n let value = value.parse().with_context(|| \"Invalid value\")?;\n config.write().save = value;\n }\n \"highlight\" => {\n let value = value.parse().with_context(|| \"Invalid value\")?;\n config.write().highlight = value;\n }\n _ => bail!(\"Unknown key '{key}'\"),\n }\n Ok(())\n }\n\n pub fn delete(config: &GlobalConfig, kind: &str) -> Result<()> {\n let (dir, file_ext) = match kind {\n \"role\" => (Self::roles_dir(), Some(\".md\")),\n \"session\" => (config.read().sessions_dir(), Some(\".yaml\")),\n \"rag\" => (Self::rags_dir(), Some(\".yaml\")),\n \"macro\" => (Self::macros_dir(), Some(\".yaml\")),\n \"agent-data\" => (Self::agents_data_dir(), None),\n _ => bail!(\"Unknown kind '{kind}'\"),\n };\n let names = match read_dir(&dir) {\n Ok(rd) => {\n let mut names = vec![];\n for entry in rd.flatten() {\n let name = entry.file_name();\n match file_ext {\n Some(file_ext) => {\n if let Some(name) = name.to_string_lossy().strip_suffix(file_ext) {\n names.push(name.to_string());\n }\n }\n None => {\n if entry.path().is_dir() {\n names.push(name.to_string_lossy().to_string());\n }\n }\n }\n }\n names.sort_unstable();\n names\n }\n Err(_) => vec![],\n };\n\n if names.is_empty() {\n bail!(\"No {kind} to delete\")\n }\n\n let select_names = MultiSelect::new(&format!(\"Select {kind} to delete:\"), names)\n .with_validator(|list: &[ListOption<&String>]| {\n if list.is_empty() {\n Ok(Validation::Invalid(\n \"At least one item must be selected\".into(),\n ))\n } else {\n Ok(Validation::Valid)\n }\n })\n .prompt()?;\n\n for name in select_names {\n match file_ext {\n Some(ext) => {\n let path = dir.join(format!(\"{name}{ext}\"));\n remove_file(&path).with_context(|| {\n format!(\"Failed to delete {kind} at '{}'\", path.display())\n })?;\n }\n None => {\n let path = dir.join(name);\n remove_dir_all(&path).with_context(|| {\n format!(\"Failed to delete {kind} at '{}'\", path.display())\n })?;\n }\n }\n }\n println!(\"✓ Successfully deleted {kind}.\");\n Ok(())\n }\n\n pub fn set_temperature(&mut self, value: Option) {\n match self.role_like_mut() {\n Some(role_like) => role_like.set_temperature(value),\n None => self.temperature = value,\n }\n }\n\n pub fn set_top_p(&mut self, value: Option) {\n match self.role_like_mut() {\n Some(role_like) => role_like.set_top_p(value),\n None => self.top_p = value,\n }\n }\n\n pub fn set_use_tools(&mut self, value: Option) {\n match self.role_like_mut() {\n Some(role_like) => role_like.set_use_tools(value),\n None => self.use_tools = value,\n }\n }\n\n pub fn set_save_session(&mut self, value: Option) {\n if let Some(session) = self.session.as_mut() {\n session.set_save_session(value);\n } else {\n self.save_session = value;\n }\n }\n\n pub fn set_compress_threshold(&mut self, value: Option) {\n if let Some(session) = self.session.as_mut() {\n session.set_compress_threshold(value);\n } else {\n self.compress_threshold = value.unwrap_or_default();\n }\n }\n\n pub fn set_rag_reranker_model(config: &GlobalConfig, value: Option) -> Result<()> {\n if let Some(id) = &value {\n Model::retrieve_model(&config.read(), id, ModelType::Reranker)?;\n }\n let has_rag = config.read().rag.is_some();\n match has_rag {\n true => update_rag(config, |rag| {\n rag.set_reranker_model(value)?;\n Ok(())\n })?,\n false => config.write().rag_reranker_model = value,\n }\n Ok(())\n }\n\n pub fn set_rag_top_k(config: &GlobalConfig, value: usize) -> Result<()> {\n let has_rag = config.read().rag.is_some();\n match has_rag {\n true => update_rag(config, |rag| {\n rag.set_top_k(value)?;\n Ok(())\n })?,\n false => config.write().rag_top_k = value,\n }\n Ok(())\n }\n\n pub fn set_wrap(&mut self, value: &str) -> Result<()> {\n if value == \"no\" {\n self.wrap = None;\n } else if value == \"auto\" {\n self.wrap = Some(value.into());\n } else {\n value\n .parse::()\n .map_err(|_| anyhow!(\"Invalid wrap value\"))?;\n self.wrap = Some(value.into())\n }\n Ok(())\n }\n\n pub fn set_max_output_tokens(&mut self, value: Option) {\n match self.role_like_mut() {\n Some(role_like) => {\n let mut model = role_like.model().clone();\n model.set_max_tokens(value, true);\n role_like.set_model(model);\n }\n None => {\n self.model.set_max_tokens(value, true);\n }\n };\n }\n\n pub fn set_model(&mut self, model_id: &str) -> Result<()> {\n let model = Model::retrieve_model(self, model_id, ModelType::Chat)?;\n match self.role_like_mut() {\n Some(role_like) => role_like.set_model(model),\n None => {\n self.model = model;\n }\n }\n Ok(())\n }\n\n pub fn use_prompt(&mut self, prompt: &str) -> Result<()> {\n let mut role = Role::new(TEMP_ROLE_NAME, prompt);\n role.set_model(self.current_model().clone());\n self.use_role_obj(role)\n }\n\n pub fn use_role(&mut self, name: &str) -> Result<()> {\n let role = self.retrieve_role(name)?;\n self.use_role_obj(role)\n }\n\n pub fn use_role_obj(&mut self, role: Role) -> Result<()> {\n if self.agent.is_some() {\n bail!(\"Cannot perform this operation because you are using a agent\")\n }\n if let Some(session) = self.session.as_mut() {\n session.guard_empty()?;\n session.set_role(role);\n } else {\n self.role = Some(role);\n }\n Ok(())\n }\n\n pub fn role_info(&self) -> Result {\n if let Some(session) = &self.session {\n if session.role_name().is_some() {\n let role = session.to_role();\n Ok(role.export())\n } else {\n bail!(\"No session role\")\n }\n } else if let Some(role) = &self.role {\n Ok(role.export())\n } else {\n bail!(\"No role\")\n }\n }\n\n pub fn exit_role(&mut self) -> Result<()> {\n if let Some(session) = self.session.as_mut() {\n session.guard_empty()?;\n session.clear_role();\n } else if self.role.is_some() {\n self.role = None;\n }\n Ok(())\n }\n\n pub fn retrieve_role(&self, name: &str) -> Result {\n let names = Self::list_roles(false);\n let mut role = if names.contains(&name.to_string()) {\n let path = Self::role_file(name);\n let content = read_to_string(&path)?;\n Role::new(name, &content)\n } else {\n Role::builtin(name)?\n };\n let current_model = self.current_model().clone();\n match role.model_id() {\n Some(model_id) => {\n if current_model.id() != model_id {\n let model = Model::retrieve_model(self, model_id, ModelType::Chat)?;\n role.set_model(model);\n } else {\n role.set_model(current_model);\n }\n }\n None => {\n role.set_model(current_model);\n if role.temperature().is_none() {\n role.set_temperature(self.temperature);\n }\n if role.top_p().is_none() {\n role.set_top_p(self.top_p);\n }\n }\n }\n Ok(role)\n }\n\n pub fn new_role(&mut self, name: &str) -> Result<()> {\n if self.macro_flag {\n bail!(\"No role\");\n }\n let ans = Confirm::new(\"Create a new role?\")\n .with_default(true)\n .prompt()?;\n if ans {\n self.upsert_role(name)?;\n } else {\n bail!(\"No role\");\n }\n Ok(())\n }\n\n pub fn edit_role(&mut self) -> Result<()> {\n let role_name;\n if let Some(session) = self.session.as_ref() {\n if let Some(name) = session.role_name().map(|v| v.to_string()) {\n if session.is_empty() {\n role_name = Some(name);\n } else {\n bail!(\"Cannot perform this operation because you are in a non-empty session\")\n }\n } else {\n bail!(\"No role\")\n }\n } else {\n role_name = self.role.as_ref().map(|v| v.name().to_string());\n }\n let name = role_name.ok_or_else(|| anyhow!(\"No role\"))?;\n self.upsert_role(&name)?;\n self.use_role(&name)\n }\n\n pub fn upsert_role(&mut self, name: &str) -> Result<()> {\n let role_path = Self::role_file(name);\n ensure_parent_exists(&role_path)?;\n let editor = self.editor()?;\n edit_file(&editor, &role_path)?;\n if self.working_mode.is_repl() {\n println!(\"✓ Saved the role to '{}'.\", role_path.display());\n }\n Ok(())\n }\n\n pub fn save_role(&mut self, name: Option<&str>) -> Result<()> {\n let mut role_name = match &self.role {\n Some(role) => {\n if role.has_args() {\n bail!(\"Unable to save the role with arguments (whose name contains '#')\")\n }\n match name {\n Some(v) => v.to_string(),\n None => role.name().to_string(),\n }\n }\n None => bail!(\"No role\"),\n };\n if role_name == TEMP_ROLE_NAME {\n role_name = Text::new(\"Role name:\")\n .with_validator(|input: &str| {\n let input = input.trim();\n if input.is_empty() {\n Ok(Validation::Invalid(\"This name is required\".into()))\n } else if input == TEMP_ROLE_NAME {\n Ok(Validation::Invalid(\"This name is reserved\".into()))\n } else {\n Ok(Validation::Valid)\n }\n })\n .prompt()?;\n }\n let role_path = Self::role_file(&role_name);\n if let Some(role) = self.role.as_mut() {\n role.save(&role_name, &role_path, self.working_mode.is_repl())?;\n }\n\n Ok(())\n }\n\n pub fn all_roles() -> Vec {\n let mut roles: HashMap = Role::list_builtin_roles()\n .iter()\n .map(|v| (v.name().to_string(), v.clone()))\n .collect();\n let names = Self::list_roles(false);\n for name in names {\n if let Ok(content) = read_to_string(Self::role_file(&name)) {\n let role = Role::new(&name, &content);\n roles.insert(name, role);\n }\n }\n let mut roles: Vec<_> = roles.into_values().collect();\n roles.sort_unstable_by(|a, b| a.name().cmp(b.name()));\n roles\n }\n\n pub fn list_roles(with_builtin: bool) -> Vec {\n let mut names = HashSet::new();\n if let Ok(rd) = read_dir(Self::roles_dir()) {\n for entry in rd.flatten() {\n if let Some(name) = entry\n .file_name()\n .to_str()\n .and_then(|v| v.strip_suffix(\".md\"))\n {\n names.insert(name.to_string());\n }\n }\n }\n if with_builtin {\n names.extend(Role::list_builtin_role_names());\n }\n let mut names: Vec<_> = names.into_iter().collect();\n names.sort_unstable();\n names\n }\n\n pub fn has_role(name: &str) -> bool {\n let names = Self::list_roles(true);\n names.contains(&name.to_string())\n }\n\n pub fn use_session(&mut self, session_name: Option<&str>) -> Result<()> {\n if self.session.is_some() {\n bail!(\n \"Already in a session, please run '.exit session' first to exit the current session.\"\n );\n }\n let mut session;\n match session_name {\n None | Some(TEMP_SESSION_NAME) => {\n let session_file = self.session_file(TEMP_SESSION_NAME);\n if session_file.exists() {\n remove_file(session_file).with_context(|| {\n format!(\"Failed to cleanup previous '{TEMP_SESSION_NAME}' session\")\n })?;\n }\n session = Some(Session::new(self, TEMP_SESSION_NAME));\n }\n Some(name) => {\n let session_path = self.session_file(name);\n if !session_path.exists() {\n session = Some(Session::new(self, name));\n } else {\n session = Some(Session::load(self, name, &session_path)?);\n }\n }\n }\n let mut new_session = false;\n if let Some(session) = session.as_mut() {\n if session.is_empty() {\n new_session = true;\n if let Some(LastMessage {\n input,\n output,\n continuous,\n }) = &self.last_message\n {\n if (*continuous && !output.is_empty())\n && self.agent.is_some() == input.with_agent()\n {\n let ans = Confirm::new(\n \"Start a session that incorporates the last question and answer?\",\n )\n .with_default(false)\n .prompt()?;\n if ans {\n session.add_message(input, output)?;\n }\n }\n }\n }\n }\n self.session = session;\n self.init_agent_session_variables(new_session)?;\n Ok(())\n }\n\n pub fn session_info(&self) -> Result {\n if let Some(session) = &self.session {\n let render_options = self.render_options()?;\n let mut markdown_render = MarkdownRender::init(render_options)?;\n let agent_info: Option<(String, Vec)> = self.agent.as_ref().map(|agent| {\n let functions = agent\n .functions()\n .declarations()\n .iter()\n .filter_map(|v| if v.agent { Some(v.name.clone()) } else { None })\n .collect();\n (agent.name().to_string(), functions)\n });\n session.render(&mut markdown_render, &agent_info)\n } else {\n bail!(\"No session\")\n }\n }\n\n pub fn exit_session(&mut self) -> Result<()> {\n if let Some(mut session) = self.session.take() {\n let sessions_dir = self.sessions_dir();\n session.exit(&sessions_dir, self.working_mode.is_repl())?;\n self.discontinuous_last_message();\n }\n Ok(())\n }\n\n pub fn save_session(&mut self, name: Option<&str>) -> Result<()> {\n let session_name = match &self.session {\n Some(session) => match name {\n Some(v) => v.to_string(),\n None => session\n .autoname()\n .unwrap_or_else(|| session.name())\n .to_string(),\n },\n None => bail!(\"No session\"),\n };\n let session_path = self.session_file(&session_name);\n if let Some(session) = self.session.as_mut() {\n session.save(&session_name, &session_path, self.working_mode.is_repl())?;\n }\n Ok(())\n }\n\n pub fn edit_session(&mut self) -> Result<()> {\n let name = match &self.session {\n Some(session) => session.name().to_string(),\n None => bail!(\"No session\"),\n };\n let session_path = self.session_file(&name);\n self.save_session(Some(&name))?;\n let editor = self.editor()?;\n edit_file(&editor, &session_path).with_context(|| {\n format!(\n \"Failed to edit '{}' with '{editor}'\",\n session_path.display()\n )\n })?;\n self.session = Some(Session::load(self, &name, &session_path)?);\n self.discontinuous_last_message();\n Ok(())\n }\n\n pub fn empty_session(&mut self) -> Result<()> {\n if let Some(session) = self.session.as_mut() {\n if let Some(agent) = self.agent.as_ref() {\n session.sync_agent(agent);\n }\n session.clear_messages();\n } else {\n bail!(\"No session\")\n }\n self.discontinuous_last_message();\n Ok(())\n }\n\n pub fn set_save_session_this_time(&mut self) -> Result<()> {\n if let Some(session) = self.session.as_mut() {\n session.set_save_session_this_time();\n } else {\n bail!(\"No session\")\n }\n Ok(())\n }\n\n pub fn list_sessions(&self) -> Vec {\n list_file_names(self.sessions_dir(), \".yaml\")\n }\n\n pub fn list_autoname_sessions(&self) -> Vec {\n list_file_names(self.sessions_dir().join(\"_\"), \".yaml\")\n }\n\n pub fn maybe_compress_session(config: GlobalConfig) {\n let mut need_compress = false;\n {\n let mut config = config.write();\n let compress_threshold = config.compress_threshold;\n if let Some(session) = config.session.as_mut() {\n if session.need_compress(compress_threshold) {\n session.set_compressing(true);\n need_compress = true;\n }\n }\n };\n if !need_compress {\n return;\n }\n let color = if config.read().light_theme() {\n nu_ansi_term::Color::LightGray\n } else {\n nu_ansi_term::Color::DarkGray\n };\n print!(\n \"\\n📢 {}\\n\",\n color.italic().paint(\"Compressing the session.\"),\n );\n tokio::spawn(async move {\n if let Err(err) = Config::compress_session(&config).await {\n warn!(\"Failed to compress the session: {err}\");\n }\n if let Some(session) = config.write().session.as_mut() {\n session.set_compressing(false);\n }\n });\n }\n\n pub async fn compress_session(config: &GlobalConfig) -> Result<()> {\n match config.read().session.as_ref() {\n Some(session) => {\n if !session.has_user_messages() {\n bail!(\"No need to compress since there are no messages in the session\")\n }\n }\n None => bail!(\"No session\"),\n }\n\n let prompt = config\n .read()\n .summarize_prompt\n .clone()\n .unwrap_or_else(|| SUMMARIZE_PROMPT.into());\n let input = Input::from_str(config, &prompt, None);\n let summary = input.fetch_chat_text().await?;\n let summary_prompt = config\n .read()\n .summary_prompt\n .clone()\n .unwrap_or_else(|| SUMMARY_PROMPT.into());\n if let Some(session) = config.write().session.as_mut() {\n session.compress(format!(\"{summary_prompt}{summary}\"));\n }\n config.write().discontinuous_last_message();\n Ok(())\n }\n\n pub fn is_compressing_session(&self) -> bool {\n self.session\n .as_ref()\n .map(|v| v.compressing())\n .unwrap_or_default()\n }\n\n pub fn maybe_autoname_session(config: GlobalConfig) {\n let mut need_autoname = false;\n if let Some(session) = config.write().session.as_mut() {\n if session.need_autoname() {\n session.set_autonaming(true);\n need_autoname = true;\n }\n }\n if !need_autoname {\n return;\n }\n let color = if config.read().light_theme() {\n nu_ansi_term::Color::LightGray\n } else {\n nu_ansi_term::Color::DarkGray\n };\n print!(\"\\n📢 {}\\n\", color.italic().paint(\"Autonaming the session.\"),);\n tokio::spawn(async move {\n if let Err(err) = Config::autoname_session(&config).await {\n warn!(\"Failed to autonaming the session: {err}\");\n }\n if let Some(session) = config.write().session.as_mut() {\n session.set_autonaming(false);\n }\n });\n }\n\n pub async fn autoname_session(config: &GlobalConfig) -> Result<()> {\n let text = match config\n .read()\n .session\n .as_ref()\n .and_then(|v| v.chat_history_for_autonaming())\n {\n Some(v) => v,\n None => bail!(\"No chat history\"),\n };\n let role = config.read().retrieve_role(CREATE_TITLE_ROLE)?;\n let input = Input::from_str(config, &text, Some(role));\n let text = input.fetch_chat_text().await?;\n if let Some(session) = config.write().session.as_mut() {\n session.set_autoname(&text);\n }\n Ok(())\n }\n\n pub async fn use_rag(\n config: &GlobalConfig,\n rag: Option<&str>,\n abort_signal: AbortSignal,\n ) -> Result<()> {\n if config.read().agent.is_some() {\n bail!(\"Cannot perform this operation because you are using a agent\")\n }\n let rag = match rag {\n None => {\n let rag_path = config.read().rag_file(TEMP_RAG_NAME);\n if rag_path.exists() {\n remove_file(&rag_path).with_context(|| {\n format!(\"Failed to cleanup previous '{TEMP_RAG_NAME}' rag\")\n })?;\n }\n Rag::init(config, TEMP_RAG_NAME, &rag_path, &[], abort_signal).await?\n }\n Some(name) => {\n let rag_path = config.read().rag_file(name);\n if !rag_path.exists() {\n if config.read().working_mode.is_cmd() {\n bail!(\"Unknown RAG '{name}'\")\n }\n Rag::init(config, name, &rag_path, &[], abort_signal).await?\n } else {\n Rag::load(config, name, &rag_path)?\n }\n }\n };\n config.write().rag = Some(Arc::new(rag));\n Ok(())\n }\n\n pub async fn edit_rag_docs(config: &GlobalConfig, abort_signal: AbortSignal) -> Result<()> {\n let mut rag = match config.read().rag.clone() {\n Some(v) => v.as_ref().clone(),\n None => bail!(\"No RAG\"),\n };\n\n let document_paths = rag.document_paths();\n let temp_file = temp_file(&format!(\"-rag-{}\", rag.name()), \".txt\");\n tokio::fs::write(&temp_file, &document_paths.join(\"\\n\"))\n .await\n .with_context(|| format!(\"Failed to write to '{}'\", temp_file.display()))?;\n let editor = config.read().editor()?;\n edit_file(&editor, &temp_file)?;\n let new_document_paths = tokio::fs::read_to_string(&temp_file)\n .await\n .with_context(|| format!(\"Failed to read '{}'\", temp_file.display()))?;\n let new_document_paths = new_document_paths\n .split('\\n')\n .filter_map(|v| {\n let v = v.trim();\n if v.is_empty() {\n None\n } else {\n Some(v.to_string())\n }\n })\n .collect::>();\n if new_document_paths.is_empty() || new_document_paths == document_paths {\n bail!(\"No changes\")\n }\n rag.refresh_document_paths(&new_document_paths, false, config, abort_signal)\n .await?;\n config.write().rag = Some(Arc::new(rag));\n Ok(())\n }\n\n pub async fn rebuild_rag(config: &GlobalConfig, abort_signal: AbortSignal) -> Result<()> {\n let mut rag = match config.read().rag.clone() {\n Some(v) => v.as_ref().clone(),\n None => bail!(\"No RAG\"),\n };\n let document_paths = rag.document_paths().to_vec();\n rag.refresh_document_paths(&document_paths, true, config, abort_signal)\n .await?;\n config.write().rag = Some(Arc::new(rag));\n Ok(())\n }\n\n pub fn rag_sources(config: &GlobalConfig) -> Result {\n match config.read().rag.as_ref() {\n Some(rag) => match rag.get_last_sources() {\n Some(v) => Ok(v),\n None => bail!(\"No sources\"),\n },\n None => bail!(\"No RAG\"),\n }\n }\n\n pub fn rag_info(&self) -> Result {\n if let Some(rag) = &self.rag {\n rag.export()\n } else {\n bail!(\"No RAG\")\n }\n }\n\n pub fn exit_rag(&mut self) -> Result<()> {\n self.rag.take();\n Ok(())\n }\n\n pub async fn search_rag(\n config: &GlobalConfig,\n rag: &Rag,\n text: &str,\n abort_signal: AbortSignal,\n ) -> Result {\n let (reranker_model, top_k) = rag.get_config();\n let (embeddings, ids) = rag\n .search(text, top_k, reranker_model.as_deref(), abort_signal)\n .await?;\n let text = config.read().rag_template(&embeddings, text);\n rag.set_last_sources(&ids);\n Ok(text)\n }\n\n pub fn list_rags() -> Vec {\n match read_dir(Self::rags_dir()) {\n Ok(rd) => {\n let mut names = vec![];\n for entry in rd.flatten() {\n let name = entry.file_name();\n if let Some(name) = name.to_string_lossy().strip_suffix(\".yaml\") {\n names.push(name.to_string());\n }\n }\n names.sort_unstable();\n names\n }\n Err(_) => vec![],\n }\n }\n\n pub fn rag_template(&self, embeddings: &str, text: &str) -> String {\n if embeddings.is_empty() {\n return text.to_string();\n }\n self.rag_template\n .as_deref()\n .unwrap_or(RAG_TEMPLATE)\n .replace(\"__CONTEXT__\", embeddings)\n .replace(\"__INPUT__\", text)\n }\n\n pub async fn use_agent(\n config: &GlobalConfig,\n agent_name: &str,\n session_name: Option<&str>,\n abort_signal: AbortSignal,\n ) -> Result<()> {\n if !config.read().function_calling {\n bail!(\"Please enable function calling before using the agent.\");\n }\n if config.read().agent.is_some() {\n bail!(\"Already in a agent, please run '.exit agent' first to exit the current agent.\");\n }\n let agent = Agent::init(config, agent_name, abort_signal).await?;\n let session = session_name.map(|v| v.to_string()).or_else(|| {\n if config.read().macro_flag {\n None\n } else {\n agent.agent_prelude().map(|v| v.to_string())\n }\n });\n config.write().rag = agent.rag();\n config.write().agent = Some(agent);\n if let Some(session) = session {\n config.write().use_session(Some(&session))?;\n } else {\n config.write().init_agent_shared_variables()?;\n }\n Ok(())\n }\n\n pub fn agent_info(&self) -> Result {\n if let Some(agent) = &self.agent {\n agent.export()\n } else {\n bail!(\"No agent\")\n }\n }\n\n pub fn agent_banner(&self) -> Result {\n if let Some(agent) = &self.agent {\n Ok(agent.banner())\n } else {\n bail!(\"No agent\")\n }\n }\n\n pub fn edit_agent_config(&self) -> Result<()> {\n let agent_name = match &self.agent {\n Some(agent) => agent.name(),\n None => bail!(\"No agent\"),\n };\n let agent_config_path = Config::agent_config_file(agent_name);\n ensure_parent_exists(&agent_config_path)?;\n if !agent_config_path.exists() {\n std::fs::write(\n &agent_config_path,\n \"# see https://github.com/sigoden/aichat/blob/main/config.agent.example.yaml\\n\",\n )\n .with_context(|| format!(\"Failed to write to '{}'\", agent_config_path.display()))?;\n }\n let editor = self.editor()?;\n edit_file(&editor, &agent_config_path)?;\n println!(\n \"NOTE: Remember to reload the agent if there are changes made to '{}'\",\n agent_config_path.display()\n );\n Ok(())\n }\n\n pub fn exit_agent(&mut self) -> Result<()> {\n self.exit_session()?;\n if self.agent.take().is_some() {\n self.rag.take();\n self.discontinuous_last_message();\n }\n Ok(())\n }\n\n pub fn exit_agent_session(&mut self) -> Result<()> {\n self.exit_session()?;\n if let Some(agent) = self.agent.as_mut() {\n agent.exit_session();\n if self.working_mode.is_repl() {\n self.init_agent_shared_variables()?;\n }\n }\n Ok(())\n }\n\n pub fn list_macros() -> Vec {\n list_file_names(Self::macros_dir(), \".yaml\")\n }\n\n pub fn load_macro(name: &str) -> Result {\n let path = Self::macro_file(name);\n let err = || format!(\"Failed to load macro '{name}' at '{}'\", path.display());\n let content = read_to_string(&path).with_context(err)?;\n let value: Macro = serde_yaml::from_str(&content).with_context(err)?;\n Ok(value)\n }\n\n pub fn has_macro(name: &str) -> bool {\n let names = Self::list_macros();\n names.contains(&name.to_string())\n }\n\n pub fn new_macro(&mut self, name: &str) -> Result<()> {\n if self.macro_flag {\n bail!(\"No macro\");\n }\n let ans = Confirm::new(\"Create a new macro?\")\n .with_default(true)\n .prompt()?;\n if ans {\n let macro_path = Self::macro_file(name);\n ensure_parent_exists(¯o_path)?;\n let editor = self.editor()?;\n edit_file(&editor, ¯o_path)?;\n } else {\n bail!(\"No macro\");\n }\n Ok(())\n }\n\n pub fn apply_prelude(&mut self) -> Result<()> {\n if self.macro_flag || !self.state().is_empty() {\n return Ok(());\n }\n let prelude = match self.working_mode {\n WorkingMode::Repl => self.repl_prelude.as_ref(),\n WorkingMode::Cmd => self.cmd_prelude.as_ref(),\n WorkingMode::Serve => return Ok(()),\n };\n let prelude = match prelude {\n Some(v) => {\n if v.is_empty() {\n return Ok(());\n }\n v.to_string()\n }\n None => return Ok(()),\n };\n\n let err_msg = || format!(\"Invalid prelude '{prelude}\");\n match prelude.split_once(':') {\n Some((\"role\", name)) => {\n self.use_role(name).with_context(err_msg)?;\n }\n Some((\"session\", name)) => {\n self.use_session(Some(name)).with_context(err_msg)?;\n }\n Some((session_name, role_name)) => {\n self.use_session(Some(session_name)).with_context(err_msg)?;\n if let Some(true) = self.session.as_ref().map(|v| v.is_empty()) {\n self.use_role(role_name).with_context(err_msg)?;\n }\n }\n _ => {\n bail!(\"{}\", err_msg())\n }\n }\n Ok(())\n }\n\n pub fn select_functions(&self, role: &Role) -> Option> {\n let mut functions = vec![];\n if self.function_calling {\n if let Some(use_tools) = role.use_tools() {\n let mut tool_names: HashSet = Default::default();\n let declaration_names: HashSet = self\n .functions\n .declarations()\n .iter()\n .map(|v| v.name.to_string())\n .collect();\n if use_tools == \"all\" {\n tool_names.extend(declaration_names);\n } else {\n for item in use_tools.split(',') {\n let item = item.trim();\n if let Some(values) = self.mapping_tools.get(item) {\n tool_names.extend(\n values\n .split(',')\n .map(|v| v.to_string())\n .filter(|v| declaration_names.contains(v)),\n )\n } else if declaration_names.contains(item) {\n tool_names.insert(item.to_string());\n }\n }\n }\n functions = self\n .functions\n .declarations()\n .iter()\n .filter_map(|v| {\n if tool_names.contains(&v.name) {\n Some(v.clone())\n } else {\n None\n }\n })\n .collect();\n }\n\n if let Some(agent) = &self.agent {\n let mut agent_functions = agent.functions().declarations().to_vec();\n let tool_names: HashSet = agent_functions\n .iter()\n .filter_map(|v| {\n if v.agent {\n None\n } else {\n Some(v.name.to_string())\n }\n })\n .collect();\n agent_functions.extend(\n functions\n .into_iter()\n .filter(|v| !tool_names.contains(&v.name)),\n );\n functions = agent_functions;\n }\n };\n if functions.is_empty() {\n None\n } else {\n Some(functions)\n }\n }\n\n pub fn editor(&self) -> Result {\n EDITOR.get_or_init(move || {\n let editor = self.editor.clone()\n .or_else(|| env::var(\"VISUAL\").ok().or_else(|| env::var(\"EDITOR\").ok()))\n .unwrap_or_else(|| {\n if cfg!(windows) {\n \"notepad\".to_string()\n } else {\n \"nano\".to_string()\n }\n });\n which::which(&editor).ok().map(|_| editor)\n })\n .clone()\n .ok_or_else(|| anyhow!(\"Editor not found. Please add the `editor` configuration or set the $EDITOR or $VISUAL environment variable.\"))\n }\n\n pub fn repl_complete(\n &self,\n cmd: &str,\n args: &[&str],\n _line: &str,\n ) -> Vec<(String, Option)> {\n let mut values: Vec<(String, Option)> = vec![];\n let filter = args.last().unwrap_or(&\"\");\n if args.len() == 1 {\n values = match cmd {\n \".role\" => map_completion_values(Self::list_roles(true)),\n \".model\" => list_models(self, ModelType::Chat)\n .into_iter()\n .map(|v| (v.id(), Some(v.description())))\n .collect(),\n \".session\" => {\n if args[0].starts_with(\"_/\") {\n map_completion_values(\n self.list_autoname_sessions()\n .iter()\n .rev()\n .map(|v| format!(\"_/{v}\"))\n .collect::>(),\n )\n } else {\n map_completion_values(self.list_sessions())\n }\n }\n \".rag\" => map_completion_values(Self::list_rags()),\n \".agent\" => map_completion_values(list_agents()),\n \".macro\" => map_completion_values(Self::list_macros()),\n \".starter\" => match &self.agent {\n Some(agent) => agent\n .conversation_staters()\n .iter()\n .enumerate()\n .map(|(i, v)| ((i + 1).to_string(), Some(v.to_string())))\n .collect(),\n None => vec![],\n },\n \".set\" => {\n let mut values = vec![\n \"temperature\",\n \"top_p\",\n \"use_tools\",\n \"save_session\",\n \"compress_threshold\",\n \"rag_reranker_model\",\n \"rag_top_k\",\n \"max_output_tokens\",\n \"dry_run\",\n \"function_calling\",\n \"stream\",\n \"save\",\n \"highlight\",\n ];\n values.sort_unstable();\n values\n .into_iter()\n .map(|v| (format!(\"{v} \"), None))\n .collect()\n }\n \".delete\" => {\n map_completion_values(vec![\"role\", \"session\", \"rag\", \"macro\", \"agent-data\"])\n }\n _ => vec![],\n };\n } else if cmd == \".set\" && args.len() == 2 {\n let candidates = match args[0] {\n \"max_output_tokens\" => match self.current_model().max_output_tokens() {\n Some(v) => vec![v.to_string()],\n None => vec![],\n },\n \"dry_run\" => complete_bool(self.dry_run),\n \"stream\" => complete_bool(self.stream),\n \"save\" => complete_bool(self.save),\n \"function_calling\" => complete_bool(self.function_calling),\n \"use_tools\" => {\n let mut prefix = String::new();\n let mut ignores = HashSet::new();\n if let Some((v, _)) = args[1].rsplit_once(',') {\n ignores = v.split(',').collect();\n prefix = format!(\"{v},\");\n }\n let mut values = vec![];\n if prefix.is_empty() {\n values.push(\"all\".to_string());\n }\n values.extend(self.functions.declarations().iter().map(|v| v.name.clone()));\n values.extend(self.mapping_tools.keys().map(|v| v.to_string()));\n values\n .into_iter()\n .filter(|v| !ignores.contains(v.as_str()))\n .map(|v| format!(\"{prefix}{v}\"))\n .collect()\n }\n \"save_session\" => {\n let save_session = if let Some(session) = &self.session {\n session.save_session()\n } else {\n self.save_session\n };\n complete_option_bool(save_session)\n }\n \"rag_reranker_model\" => list_models(self, ModelType::Reranker)\n .iter()\n .map(|v| v.id())\n .collect(),\n \"highlight\" => complete_bool(self.highlight),\n _ => vec![],\n };\n values = candidates.into_iter().map(|v| (v, None)).collect();\n } else if cmd == \".agent\" {\n if args.len() == 2 {\n let dir = Self::agent_data_dir(args[0]).join(SESSIONS_DIR_NAME);\n values = list_file_names(dir, \".yaml\")\n .into_iter()\n .map(|v| (v, None))\n .collect();\n }\n values.extend(complete_agent_variables(args[0]));\n };\n fuzzy_filter(values, |v| v.0.as_str(), filter)\n }\n\n pub fn sync_models_url(&self) -> String {\n self.sync_models_url\n .clone()\n .unwrap_or_else(|| SYNC_MODELS_URL.into())\n }\n\n pub async fn sync_models(url: &str, abort_signal: AbortSignal) -> Result<()> {\n let content = abortable_run_with_spinner(fetch(url), \"Fetching models.yaml\", abort_signal)\n .await\n .with_context(|| format!(\"Failed to fetch '{url}'\"))?;\n println!(\"✓ Fetched '{url}'\");\n let list = serde_yaml::from_str::>(&content)\n .with_context(|| \"Failed to parse models.yaml\")?;\n let models_override = ModelsOverride {\n version: env!(\"CARGO_PKG_VERSION\").to_string(),\n list,\n };\n let models_override_data =\n serde_yaml::to_string(&models_override).with_context(|| \"Failed to serde {}\")?;\n\n let model_override_path = Self::models_override_file();\n ensure_parent_exists(&model_override_path)?;\n std::fs::write(&model_override_path, models_override_data)\n .with_context(|| format!(\"Failed to write to '{}'\", model_override_path.display()))?;\n println!(\"✓ Updated '{}'\", model_override_path.display());\n Ok(())\n }\n\n pub fn loal_models_override() -> Result> {\n let model_override_path = Self::models_override_file();\n let err = || {\n format!(\n \"Failed to load models at '{}'\",\n model_override_path.display()\n )\n };\n let content = read_to_string(&model_override_path).with_context(err)?;\n let models_override: ModelsOverride = serde_yaml::from_str(&content).with_context(err)?;\n if models_override.version != env!(\"CARGO_PKG_VERSION\") {\n bail!(\"Incompatible version\")\n }\n Ok(models_override.list)\n }\n\n pub fn light_theme(&self) -> bool {\n matches!(self.theme.as_deref(), Some(\"light\"))\n }\n\n pub fn render_options(&self) -> Result {\n let theme = if self.highlight {\n let theme_mode = if self.light_theme() { \"light\" } else { \"dark\" };\n let theme_filename = format!(\"{theme_mode}.tmTheme\");\n let theme_path = Self::local_path(&theme_filename);\n if theme_path.exists() {\n let theme = ThemeSet::get_theme(&theme_path)\n .with_context(|| format!(\"Invalid theme at '{}'\", theme_path.display()))?;\n Some(theme)\n } else {\n let theme = if self.light_theme() {\n decode_bin(LIGHT_THEME).context(\"Invalid builtin light theme\")?\n } else {\n decode_bin(DARK_THEME).context(\"Invalid builtin dark theme\")?\n };\n Some(theme)\n }\n } else {\n None\n };\n let wrap = if *IS_STDOUT_TERMINAL {\n self.wrap.clone()\n } else {\n None\n };\n let truecolor = matches!(\n env::var(\"COLORTERM\").as_ref().map(|v| v.as_str()),\n Ok(\"truecolor\")\n );\n Ok(RenderOptions::new(theme, wrap, self.wrap_code, truecolor))\n }\n\n pub fn render_prompt_left(&self) -> String {\n let variables = self.generate_prompt_context();\n let left_prompt = self.left_prompt.as_deref().unwrap_or(LEFT_PROMPT);\n render_prompt(left_prompt, &variables)\n }\n\n pub fn render_prompt_right(&self) -> String {\n let variables = self.generate_prompt_context();\n let right_prompt = self.right_prompt.as_deref().unwrap_or(RIGHT_PROMPT);\n render_prompt(right_prompt, &variables)\n }\n\n pub fn print_markdown(&self, text: &str) -> Result<()> {\n if *IS_STDOUT_TERMINAL {\n let render_options = self.render_options()?;\n let mut markdown_render = MarkdownRender::init(render_options)?;\n println!(\"{}\", markdown_render.render(text));\n } else {\n println!(\"{text}\");\n }\n Ok(())\n }\n\n fn generate_prompt_context(&self) -> HashMap<&str, String> {\n let mut output = HashMap::new();\n let role = self.extract_role();\n output.insert(\"model\", role.model().id());\n output.insert(\"client_name\", role.model().client_name().to_string());\n output.insert(\"model_name\", role.model().name().to_string());\n output.insert(\n \"max_input_tokens\",\n role.model()\n .max_input_tokens()\n .unwrap_or_default()\n .to_string(),\n );\n if let Some(temperature) = role.temperature() {\n if temperature != 0.0 {\n output.insert(\"temperature\", temperature.to_string());\n }\n }\n if let Some(top_p) = role.top_p() {\n if top_p != 0.0 {\n output.insert(\"top_p\", top_p.to_string());\n }\n }\n if self.dry_run {\n output.insert(\"dry_run\", \"true\".to_string());\n }\n if self.stream {\n output.insert(\"stream\", \"true\".to_string());\n }\n if self.save {\n output.insert(\"save\", \"true\".to_string());\n }\n if let Some(wrap) = &self.wrap {\n if wrap != \"no\" {\n output.insert(\"wrap\", wrap.clone());\n }\n }\n if !role.is_derived() {\n output.insert(\"role\", role.name().to_string());\n }\n if let Some(session) = &self.session {\n output.insert(\"session\", session.name().to_string());\n if let Some(autoname) = session.autoname() {\n output.insert(\"session_autoname\", autoname.to_string());\n }\n output.insert(\"dirty\", session.dirty().to_string());\n let (tokens, percent) = session.tokens_usage();\n output.insert(\"consume_tokens\", tokens.to_string());\n output.insert(\"consume_percent\", percent.to_string());\n output.insert(\"user_messages_len\", session.user_messages_len().to_string());\n }\n if let Some(rag) = &self.rag {\n output.insert(\"rag\", rag.name().to_string());\n }\n if let Some(agent) = &self.agent {\n output.insert(\"agent\", agent.name().to_string());\n }\n\n if self.highlight {\n output.insert(\"color.reset\", \"\\u{1b}[0m\".to_string());\n output.insert(\"color.black\", \"\\u{1b}[30m\".to_string());\n output.insert(\"color.dark_gray\", \"\\u{1b}[90m\".to_string());\n output.insert(\"color.red\", \"\\u{1b}[31m\".to_string());\n output.insert(\"color.light_red\", \"\\u{1b}[91m\".to_string());\n output.insert(\"color.green\", \"\\u{1b}[32m\".to_string());\n output.insert(\"color.light_green\", \"\\u{1b}[92m\".to_string());\n output.insert(\"color.yellow\", \"\\u{1b}[33m\".to_string());\n output.insert(\"color.light_yellow\", \"\\u{1b}[93m\".to_string());\n output.insert(\"color.blue\", \"\\u{1b}[34m\".to_string());\n output.insert(\"color.light_blue\", \"\\u{1b}[94m\".to_string());\n output.insert(\"color.purple\", \"\\u{1b}[35m\".to_string());\n output.insert(\"color.light_purple\", \"\\u{1b}[95m\".to_string());\n output.insert(\"color.magenta\", \"\\u{1b}[35m\".to_string());\n output.insert(\"color.light_magenta\", \"\\u{1b}[95m\".to_string());\n output.insert(\"color.cyan\", \"\\u{1b}[36m\".to_string());\n output.insert(\"color.light_cyan\", \"\\u{1b}[96m\".to_string());\n output.insert(\"color.white\", \"\\u{1b}[37m\".to_string());\n output.insert(\"color.light_gray\", \"\\u{1b}[97m\".to_string());\n }\n\n output\n }\n\n pub fn before_chat_completion(&mut self, input: &Input) -> Result<()> {\n self.last_message = Some(LastMessage::new(input.clone(), String::new()));\n Ok(())\n }\n\n pub fn after_chat_completion(\n &mut self,\n input: &Input,\n output: &str,\n tool_results: &[ToolResult],\n ) -> Result<()> {\n if !tool_results.is_empty() {\n return Ok(());\n }\n self.last_message = Some(LastMessage::new(input.clone(), output.to_string()));\n if !self.dry_run {\n self.save_message(input, output)?;\n }\n Ok(())\n }\n\n fn discontinuous_last_message(&mut self) {\n if let Some(last_message) = self.last_message.as_mut() {\n last_message.continuous = false;\n }\n }\n\n fn save_message(&mut self, input: &Input, output: &str) -> Result<()> {\n let mut input = input.clone();\n input.clear_patch();\n if let Some(session) = input.session_mut(&mut self.session) {\n session.add_message(&input, output)?;\n return Ok(());\n }\n\n if !self.save {\n return Ok(());\n }\n let mut file = self.open_message_file()?;\n if output.is_empty() && input.tool_calls().is_none() {\n return Ok(());\n }\n let now = now();\n let summary = input.summary();\n let raw_input = input.raw();\n let scope = if self.agent.is_none() {\n let role_name = if input.role().is_derived() {\n None\n } else {\n Some(input.role().name())\n };\n match (role_name, input.rag_name()) {\n (Some(role), Some(rag_name)) => format!(\" ({role}#{rag_name})\"),\n (Some(role), _) => format!(\" ({role})\"),\n (None, Some(rag_name)) => format!(\" (#{rag_name})\"),\n _ => String::new(),\n }\n } else {\n String::new()\n };\n let tool_calls = match input.tool_calls() {\n Some(MessageContentToolCalls {\n tool_results, text, ..\n }) => {\n let mut lines = vec![\"\".to_string()];\n if !text.is_empty() {\n lines.push(text.clone());\n }\n lines.push(serde_json::to_string(&tool_results).unwrap_or_default());\n lines.push(\"\\n\".to_string());\n lines.join(\"\\n\")\n }\n None => String::new(),\n };\n let output = format!(\n \"# CHAT: {summary} [{now}]{scope}\\n{raw_input}\\n--------\\n{tool_calls}{output}\\n--------\\n\\n\",\n );\n file.write_all(output.as_bytes())\n .with_context(|| \"Failed to save message\")\n }\n\n fn init_agent_shared_variables(&mut self) -> Result<()> {\n let agent = match self.agent.as_mut() {\n Some(v) => v,\n None => return Ok(()),\n };\n if !agent.defined_variables().is_empty() && agent.shared_variables().is_empty() {\n let mut config_variables = agent.config_variables().clone();\n if let Some(v) = &self.agent_variables {\n config_variables.extend(v.clone());\n }\n let new_variables = Agent::init_agent_variables(\n agent.defined_variables(),\n &config_variables,\n self.info_flag,\n )?;\n agent.set_shared_variables(new_variables);\n }\n if !self.info_flag {\n agent.update_shared_dynamic_instructions(false)?;\n }\n Ok(())\n }\n\n fn init_agent_session_variables(&mut self, new_session: bool) -> Result<()> {\n let (agent, session) = match (self.agent.as_mut(), self.session.as_mut()) {\n (Some(agent), Some(session)) => (agent, session),\n _ => return Ok(()),\n };\n if new_session {\n let shared_variables = agent.shared_variables().clone();\n let session_variables =\n if !agent.defined_variables().is_empty() && shared_variables.is_empty() {\n let mut config_variables = agent.config_variables().clone();\n if let Some(v) = &self.agent_variables {\n config_variables.extend(v.clone());\n }\n let new_variables = Agent::init_agent_variables(\n agent.defined_variables(),\n &config_variables,\n self.info_flag,\n )?;\n agent.set_shared_variables(new_variables.clone());\n new_variables\n } else {\n shared_variables\n };\n agent.set_session_variables(session_variables);\n if !self.info_flag {\n agent.update_session_dynamic_instructions(None)?;\n }\n session.sync_agent(agent);\n } else {\n let variables = session.agent_variables();\n agent.set_session_variables(variables.clone());\n agent.update_session_dynamic_instructions(Some(\n session.agent_instructions().to_string(),\n ))?;\n }\n Ok(())\n }\n\n fn open_message_file(&self) -> Result {\n let path = self.messages_file();\n ensure_parent_exists(&path)?;\n OpenOptions::new()\n .create(true)\n .append(true)\n .open(&path)\n .with_context(|| format!(\"Failed to create/append {}\", path.display()))\n }\n\n fn load_from_file(config_path: &Path) -> Result {\n let err = || format!(\"Failed to load config at '{}'\", config_path.display());\n let content = read_to_string(config_path).with_context(err)?;\n let config: Self = serde_yaml::from_str(&content)\n .map_err(|err| {\n let err_msg = err.to_string();\n let err_msg = if err_msg.starts_with(&format!(\"{CLIENTS_FIELD}: \")) {\n // location is incorrect, get rid of it\n err_msg\n .split_once(\" at line\")\n .map(|(v, _)| {\n format!(\"{v} (Sorry for being unable to provide an exact location)\")\n })\n .unwrap_or_else(|| \"clients: invalid value\".into())\n } else {\n err_msg\n };\n anyhow!(\"{err_msg}\")\n })\n .with_context(err)?;\n\n Ok(config)\n }\n\n fn load_dynamic(model_id: &str) -> Result {\n let provider = match model_id.split_once(':') {\n Some((v, _)) => v,\n _ => model_id,\n };\n let is_openai_compatible = OPENAI_COMPATIBLE_PROVIDERS\n .into_iter()\n .any(|(name, _)| provider == name);\n let client = if is_openai_compatible {\n json!({ \"type\": \"openai-compatible\", \"name\": provider })\n } else {\n json!({ \"type\": provider })\n };\n let config = json!({\n \"model\": model_id.to_string(),\n \"save\": false,\n \"clients\": vec![client],\n });\n let config =\n serde_json::from_value(config).with_context(|| \"Failed to load config from env\")?;\n Ok(config)\n }\n\n fn load_envs(&mut self) {\n if let Ok(v) = env::var(get_env_name(\"model\")) {\n self.model_id = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"temperature\")) {\n self.temperature = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"top_p\")) {\n self.top_p = v;\n }\n\n if let Some(Some(v)) = read_env_bool(&get_env_name(\"dry_run\")) {\n self.dry_run = v;\n }\n if let Some(Some(v)) = read_env_bool(&get_env_name(\"stream\")) {\n self.stream = v;\n }\n if let Some(Some(v)) = read_env_bool(&get_env_name(\"save\")) {\n self.save = v;\n }\n if let Ok(v) = env::var(get_env_name(\"keybindings\")) {\n if v == \"vi\" {\n self.keybindings = v;\n }\n }\n if let Some(v) = read_env_value::(&get_env_name(\"editor\")) {\n self.editor = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"wrap\")) {\n self.wrap = v;\n }\n if let Some(Some(v)) = read_env_bool(&get_env_name(\"wrap_code\")) {\n self.wrap_code = v;\n }\n\n if let Some(Some(v)) = read_env_bool(&get_env_name(\"function_calling\")) {\n self.function_calling = v;\n }\n if let Ok(v) = env::var(get_env_name(\"mapping_tools\")) {\n if let Ok(v) = serde_json::from_str(&v) {\n self.mapping_tools = v;\n }\n }\n if let Some(v) = read_env_value::(&get_env_name(\"use_tools\")) {\n self.use_tools = v;\n }\n\n if let Some(v) = read_env_value::(&get_env_name(\"repl_prelude\")) {\n self.repl_prelude = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"cmd_prelude\")) {\n self.cmd_prelude = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"agent_prelude\")) {\n self.agent_prelude = v;\n }\n\n if let Some(v) = read_env_bool(&get_env_name(\"save_session\")) {\n self.save_session = v;\n }\n if let Some(Some(v)) = read_env_value::(&get_env_name(\"compress_threshold\")) {\n self.compress_threshold = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"summarize_prompt\")) {\n self.summarize_prompt = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"summary_prompt\")) {\n self.summary_prompt = v;\n }\n\n if let Some(v) = read_env_value::(&get_env_name(\"rag_embedding_model\")) {\n self.rag_embedding_model = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"rag_reranker_model\")) {\n self.rag_reranker_model = v;\n }\n if let Some(Some(v)) = read_env_value::(&get_env_name(\"rag_top_k\")) {\n self.rag_top_k = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"rag_chunk_size\")) {\n self.rag_chunk_size = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"rag_chunk_overlap\")) {\n self.rag_chunk_overlap = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"rag_template\")) {\n self.rag_template = v;\n }\n\n if let Ok(v) = env::var(get_env_name(\"document_loaders\")) {\n if let Ok(v) = serde_json::from_str(&v) {\n self.document_loaders = v;\n }\n }\n\n if let Some(Some(v)) = read_env_bool(&get_env_name(\"highlight\")) {\n self.highlight = v;\n }\n if *NO_COLOR {\n self.highlight = false;\n }\n if self.highlight && self.theme.is_none() {\n if let Some(v) = read_env_value::(&get_env_name(\"theme\")) {\n self.theme = v;\n } else if *IS_STDOUT_TERMINAL {\n if let Ok(color_scheme) = color_scheme(QueryOptions::default()) {\n let theme = match color_scheme {\n ColorScheme::Dark => \"dark\",\n ColorScheme::Light => \"light\",\n };\n self.theme = Some(theme.into());\n }\n }\n }\n if let Some(v) = read_env_value::(&get_env_name(\"left_prompt\")) {\n self.left_prompt = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"right_prompt\")) {\n self.right_prompt = v;\n }\n\n if let Some(v) = read_env_value::(&get_env_name(\"serve_addr\")) {\n self.serve_addr = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"user_agent\")) {\n self.user_agent = v;\n }\n if let Some(Some(v)) = read_env_bool(&get_env_name(\"save_shell_history\")) {\n self.save_shell_history = v;\n }\n if let Some(v) = read_env_value::(&get_env_name(\"sync_models_url\")) {\n self.sync_models_url = v;\n }\n }\n\n fn load_functions(&mut self) -> Result<()> {\n self.functions = Functions::init(&Self::functions_file())?;\n Ok(())\n }\n\n fn setup_model(&mut self) -> Result<()> {\n let mut model_id = self.model_id.clone();\n if model_id.is_empty() {\n let models = list_models(self, ModelType::Chat);\n if models.is_empty() {\n bail!(\"No available model\");\n }\n model_id = models[0].id()\n };\n self.set_model(&model_id)?;\n self.model_id = model_id;\n Ok(())\n }\n\n fn setup_document_loaders(&mut self) {\n [(\"pdf\", \"pdftotext $1 -\"), (\"docx\", \"pandoc --to plain $1\")]\n .into_iter()\n .for_each(|(k, v)| {\n let (k, v) = (k.to_string(), v.to_string());\n self.document_loaders.entry(k).or_insert(v);\n });\n }\n\n fn setup_user_agent(&mut self) {\n if let Some(\"auto\") = self.user_agent.as_deref() {\n self.user_agent = Some(format!(\n \"{}/{}\",\n env!(\"CARGO_CRATE_NAME\"),\n env!(\"CARGO_PKG_VERSION\")\n ));\n }\n }\n}\n\npub fn load_env_file() -> Result<()> {\n let env_file_path = Config::env_file();\n let contents = match read_to_string(&env_file_path) {\n Ok(v) => v,\n Err(_) => return Ok(()),\n };\n debug!(\"Use env file '{}'\", env_file_path.display());\n for line in contents.lines() {\n let line = line.trim();\n if line.starts_with('#') || line.is_empty() {\n continue;\n }\n if let Some((key, value)) = line.split_once('=') {\n env::set_var(key.trim(), value.trim());\n }\n }\n Ok(())\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]\npub enum WorkingMode {\n Cmd,\n Repl,\n Serve,\n}\n\nimpl WorkingMode {\n pub fn is_cmd(&self) -> bool {\n *self == WorkingMode::Cmd\n }\n pub fn is_repl(&self) -> bool {\n *self == WorkingMode::Repl\n }\n pub fn is_serve(&self) -> bool {\n *self == WorkingMode::Serve\n }\n}\n\n#[async_recursion::async_recursion]\npub async fn macro_execute(\n config: &GlobalConfig,\n name: &str,\n args: Option<&str>,\n abort_signal: AbortSignal,\n) -> Result<()> {\n let macro_value = Config::load_macro(name)?;\n let (mut new_args, text) = split_args_text(args.unwrap_or_default(), cfg!(windows));\n if !text.is_empty() {\n new_args.push(text.to_string());\n }\n let variables = macro_value\n .resolve_variables(&new_args)\n .map_err(|err| anyhow!(\"{err}. Usage: {}\", macro_value.usage(name)))?;\n let role = config.read().extract_role();\n let mut config = config.read().clone();\n config.temperature = role.temperature();\n config.top_p = role.top_p();\n config.use_tools = role.use_tools().clone();\n config.macro_flag = true;\n config.model = role.model().clone();\n config.role = None;\n config.session = None;\n config.rag = None;\n config.agent = None;\n config.discontinuous_last_message();\n let config = Arc::new(RwLock::new(config));\n config.write().macro_flag = true;\n for step in ¯o_value.steps {\n let command = Macro::interpolate_command(step, &variables);\n println!(\">> {}\", multiline_text(&command));\n run_repl_command(&config, abort_signal.clone(), &command).await?;\n }\n Ok(())\n}\n\n#[derive(Debug, Clone, Deserialize)]\npub struct Macro {\n #[serde(default)]\n pub variables: Vec,\n pub steps: Vec,\n}\n\nimpl Macro {\n pub fn resolve_variables(&self, args: &[String]) -> Result> {\n let mut output = IndexMap::new();\n for (i, variable) in self.variables.iter().enumerate() {\n let value = if variable.rest && i == self.variables.len() - 1 {\n if args.len() > i {\n Some(args[i..].join(\" \"))\n } else {\n variable.default.clone()\n }\n } else {\n args.get(i)\n .map(|v| v.to_string())\n .or_else(|| variable.default.clone())\n };\n let value =\n value.ok_or_else(|| anyhow!(\"Missing value for variable '{}'\", variable.name))?;\n output.insert(variable.name.clone(), value);\n }\n Ok(output)\n }\n\n pub fn usage(&self, name: &str) -> String {\n let mut parts = vec![name.to_string()];\n for (i, variable) in self.variables.iter().enumerate() {\n let part = match (\n variable.rest && i == self.variables.len() - 1,\n variable.default.is_some(),\n ) {\n (true, true) => format!(\"[{}]...\", variable.name),\n (true, false) => format!(\"<{}>...\", variable.name),\n (false, true) => format!(\"[{}]\", variable.name),\n (false, false) => format!(\"<{}>\", variable.name),\n };\n parts.push(part);\n }\n parts.join(\" \")\n }\n\n pub fn interpolate_command(command: &str, variables: &IndexMap) -> String {\n let mut output = command.to_string();\n for (key, value) in variables {\n output = output.replace(&format!(\"{{{{{key}}}}}\"), value);\n }\n output\n }\n}\n\n#[derive(Debug, Clone, Deserialize)]\npub struct MacroVariable {\n pub name: String,\n #[serde(default)]\n pub rest: bool,\n pub default: Option,\n}\n\n#[derive(Debug, Clone, Serialize, Deserialize)]\npub struct ModelsOverride {\n pub version: String,\n pub list: Vec,\n}\n\n#[derive(Debug, Clone)]\npub struct LastMessage {\n pub input: Input,\n pub output: String,\n pub continuous: bool,\n}\n\nimpl LastMessage {\n pub fn new(input: Input, output: String) -> Self {\n Self {\n input,\n output,\n continuous: true,\n }\n }\n}\n\nbitflags::bitflags! {\n #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]\n pub struct StateFlags: u32 {\n const ROLE = 1 << 0;\n const SESSION_EMPTY = 1 << 1;\n const SESSION = 1 << 2;\n const RAG = 1 << 3;\n const AGENT = 1 << 4;\n }\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]\npub enum AssertState {\n True(StateFlags),\n False(StateFlags),\n TrueFalse(StateFlags, StateFlags),\n Equal(StateFlags),\n}\n\nimpl AssertState {\n pub fn pass() -> Self {\n AssertState::False(StateFlags::empty())\n }\n\n pub fn bare() -> Self {\n AssertState::Equal(StateFlags::empty())\n }\n\n pub fn assert(self, flags: StateFlags) -> bool {\n match self {\n AssertState::True(true_flags) => true_flags & flags != StateFlags::empty(),\n AssertState::False(false_flags) => false_flags & flags == StateFlags::empty(),\n AssertState::TrueFalse(true_flags, false_flags) => {\n (true_flags & flags != StateFlags::empty())\n && (false_flags & flags == StateFlags::empty())\n }\n AssertState::Equal(check_flags) => check_flags == flags,\n }\n }\n}\n\nasync fn create_config_file(config_path: &Path) -> Result<()> {\n let ans = Confirm::new(\"No config file, create a new one?\")\n .with_default(true)\n .prompt()?;\n if !ans {\n process::exit(0);\n }\n\n let client = Select::new(\"API Provider (required):\", list_client_types()).prompt()?;\n\n let mut config = serde_json::json!({});\n let (model, clients_config) = create_client_config(client).await?;\n config[\"model\"] = model.into();\n config[CLIENTS_FIELD] = clients_config;\n\n let config_data = serde_yaml::to_string(&config).with_context(|| \"Failed to create config\")?;\n let config_data = format!(\n \"# see https://github.com/sigoden/aichat/blob/main/config.example.yaml\\n\\n{config_data}\"\n );\n\n ensure_parent_exists(config_path)?;\n std::fs::write(config_path, config_data)\n .with_context(|| format!(\"Failed to write to '{}'\", config_path.display()))?;\n #[cfg(unix)]\n {\n use std::os::unix::prelude::PermissionsExt;\n let perms = std::fs::Permissions::from_mode(0o600);\n std::fs::set_permissions(config_path, perms)?;\n }\n\n println!(\"✓ Saved the config file to '{}'.\\n\", config_path.display());\n\n Ok(())\n}\n\npub(crate) fn ensure_parent_exists(path: &Path) -> Result<()> {\n if path.exists() {\n return Ok(());\n }\n let parent = path\n .parent()\n .ok_or_else(|| anyhow!(\"Failed to write to '{}', No parent path\", path.display()))?;\n if !parent.exists() {\n create_dir_all(parent).with_context(|| {\n format!(\n \"Failed to write to '{}', Cannot create parent directory\",\n path.display()\n )\n })?;\n }\n Ok(())\n}\n\nfn read_env_value(key: &str) -> Option>\nwhere\n T: std::str::FromStr,\n{\n let value = env::var(key).ok()?;\n let value = parse_value(&value).ok()?;\n Some(value)\n}\n\nfn parse_value(value: &str) -> Result>\nwhere\n T: std::str::FromStr,\n{\n let value = if value == \"null\" {\n None\n } else {\n let value = match value.parse() {\n Ok(value) => value,\n Err(_) => bail!(\"Invalid value '{}'\", value),\n };\n Some(value)\n };\n Ok(value)\n}\n\nfn read_env_bool(key: &str) -> Option> {\n let value = env::var(key).ok()?;\n Some(parse_bool(&value))\n}\n\nfn complete_bool(value: bool) -> Vec {\n vec![(!value).to_string()]\n}\n\nfn complete_option_bool(value: Option) -> Vec {\n match value {\n Some(true) => vec![\"false\".to_string(), \"null\".to_string()],\n Some(false) => vec![\"true\".to_string(), \"null\".to_string()],\n None => vec![\"true\".to_string(), \"false\".to_string()],\n }\n}\n\nfn map_completion_values(value: Vec) -> Vec<(String, Option)> {\n value.into_iter().map(|v| (v.to_string(), None)).collect()\n}\n\nfn update_rag(config: &GlobalConfig, f: F) -> Result<()>\nwhere\n F: FnOnce(&mut Rag) -> Result<()>,\n{\n let mut rag = match config.read().rag.clone() {\n Some(v) => v.as_ref().clone(),\n None => bail!(\"No RAG\"),\n };\n f(&mut rag)?;\n config.write().rag = Some(Arc::new(rag));\n Ok(())\n}\n\nfn format_option_value(value: &Option) -> String\nwhere\n T: std::fmt::Display,\n{\n match value {\n Some(value) => value.to_string(),\n None => \"null\".to_string(),\n }\n}\n" }, { "path": "src/config/role.rs", "content": "use super::*;\n\nuse crate::client::{Message, MessageContent, MessageRole, Model};\n\nuse anyhow::Result;\nuse fancy_regex::Regex;\nuse rust_embed::Embed;\nuse serde::{Deserialize, Serialize};\nuse serde_json::Value;\nuse std::sync::LazyLock;\n\npub const SHELL_ROLE: &str = \"%shell%\";\npub const EXPLAIN_SHELL_ROLE: &str = \"%explain-shell%\";\npub const CODE_ROLE: &str = \"%code%\";\npub const CREATE_TITLE_ROLE: &str = \"%create-title%\";\n\npub const INPUT_PLACEHOLDER: &str = \"__INPUT__\";\n\n#[derive(Embed)]\n#[folder = \"assets/roles/\"]\nstruct RolesAsset;\n\nstatic RE_METADATA: LazyLock =\n LazyLock::new(|| Regex::new(r\"(?s)-{3,}\\s*(.*?)\\s*-{3,}\\s*(.*)\").unwrap());\n\npub trait RoleLike {\n fn to_role(&self) -> Role;\n fn model(&self) -> &Model;\n fn temperature(&self) -> Option;\n fn top_p(&self) -> Option;\n fn use_tools(&self) -> Option;\n fn set_model(&mut self, model: Model);\n fn set_temperature(&mut self, value: Option);\n fn set_top_p(&mut self, value: Option);\n fn set_use_tools(&mut self, value: Option);\n}\n\n#[derive(Debug, Clone, Default, Deserialize, Serialize)]\npub struct Role {\n name: String,\n #[serde(default)]\n prompt: String,\n #[serde(\n rename(serialize = \"model\", deserialize = \"model\"),\n skip_serializing_if = \"Option::is_none\"\n )]\n model_id: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n temperature: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n top_p: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n use_tools: Option,\n\n #[serde(skip)]\n model: Model,\n}\n\nimpl Role {\n pub fn new(name: &str, content: &str) -> Self {\n let mut metadata = \"\";\n let mut prompt = content.trim();\n if let Ok(Some(caps)) = RE_METADATA.captures(content) {\n if let (Some(metadata_value), Some(prompt_value)) = (caps.get(1), caps.get(2)) {\n metadata = metadata_value.as_str().trim();\n prompt = prompt_value.as_str().trim();\n }\n }\n let mut prompt = prompt.to_string();\n interpolate_variables(&mut prompt);\n let mut role = Self {\n name: name.to_string(),\n prompt,\n ..Default::default()\n };\n if !metadata.is_empty() {\n if let Ok(value) = serde_yaml::from_str::(metadata) {\n if let Some(value) = value.as_object() {\n for (key, value) in value {\n match key.as_str() {\n \"model\" => role.model_id = value.as_str().map(|v| v.to_string()),\n \"temperature\" => role.temperature = value.as_f64(),\n \"top_p\" => role.top_p = value.as_f64(),\n \"use_tools\" => role.use_tools = value.as_str().map(|v| v.to_string()),\n _ => (),\n }\n }\n }\n }\n }\n role\n }\n\n pub fn builtin(name: &str) -> Result {\n let content = RolesAsset::get(&format!(\"{name}.md\"))\n .ok_or_else(|| anyhow!(\"Unknown role `{name}`\"))?;\n let content = unsafe { std::str::from_utf8_unchecked(&content.data) };\n Ok(Role::new(name, content))\n }\n\n pub fn list_builtin_role_names() -> Vec {\n RolesAsset::iter()\n .filter_map(|v| v.strip_suffix(\".md\").map(|v| v.to_string()))\n .collect()\n }\n\n pub fn list_builtin_roles() -> Vec {\n RolesAsset::iter()\n .filter_map(|v| Role::builtin(&v).ok())\n .collect()\n }\n\n pub fn has_args(&self) -> bool {\n self.name.contains('#')\n }\n\n pub fn export(&self) -> String {\n let mut metadata = vec![];\n if let Some(model) = self.model_id() {\n metadata.push(format!(\"model: {model}\"));\n }\n if let Some(temperature) = self.temperature() {\n metadata.push(format!(\"temperature: {temperature}\"));\n }\n if let Some(top_p) = self.top_p() {\n metadata.push(format!(\"top_p: {top_p}\"));\n }\n if let Some(use_tools) = self.use_tools() {\n metadata.push(format!(\"use_tools: {use_tools}\"));\n }\n if metadata.is_empty() {\n format!(\"{}\\n\", self.prompt)\n } else if self.prompt.is_empty() {\n format!(\"---\\n{}\\n---\\n\", metadata.join(\"\\n\"))\n } else {\n format!(\"---\\n{}\\n---\\n\\n{}\\n\", metadata.join(\"\\n\"), self.prompt)\n }\n }\n\n pub fn save(&mut self, role_name: &str, role_path: &Path, is_repl: bool) -> Result<()> {\n ensure_parent_exists(role_path)?;\n\n let content = self.export();\n std::fs::write(role_path, content).with_context(|| {\n format!(\n \"Failed to write role {} to {}\",\n self.name,\n role_path.display()\n )\n })?;\n\n if is_repl {\n println!(\"✓ Saved role to '{}'.\", role_path.display());\n }\n\n if role_name != self.name {\n self.name = role_name.to_string();\n }\n\n Ok(())\n }\n\n pub fn sync(&mut self, role_like: &T) {\n let model = role_like.model();\n let temperature = role_like.temperature();\n let top_p = role_like.top_p();\n let use_tools = role_like.use_tools();\n self.batch_set(model, temperature, top_p, use_tools);\n }\n\n pub fn batch_set(\n &mut self,\n model: &Model,\n temperature: Option,\n top_p: Option,\n use_tools: Option,\n ) {\n self.set_model(model.clone());\n if temperature.is_some() {\n self.set_temperature(temperature);\n }\n if top_p.is_some() {\n self.set_top_p(top_p);\n }\n if use_tools.is_some() {\n self.set_use_tools(use_tools);\n }\n }\n\n pub fn is_derived(&self) -> bool {\n self.name.is_empty()\n }\n\n pub fn name(&self) -> &str {\n &self.name\n }\n\n pub fn model_id(&self) -> Option<&str> {\n self.model_id.as_deref()\n }\n\n pub fn prompt(&self) -> &str {\n &self.prompt\n }\n\n pub fn is_empty_prompt(&self) -> bool {\n self.prompt.is_empty()\n }\n\n pub fn is_embedded_prompt(&self) -> bool {\n self.prompt.contains(INPUT_PLACEHOLDER)\n }\n\n pub fn echo_messages(&self, input: &Input) -> String {\n let input_markdown = input.render();\n if self.is_empty_prompt() {\n input_markdown\n } else if self.is_embedded_prompt() {\n self.prompt.replace(INPUT_PLACEHOLDER, &input_markdown)\n } else {\n format!(\"{}\\n\\n{}\", self.prompt, input_markdown)\n }\n }\n\n pub fn build_messages(&self, input: &Input) -> Vec {\n let mut content = input.message_content();\n let mut messages = if self.is_empty_prompt() {\n vec![Message::new(MessageRole::User, content)]\n } else if self.is_embedded_prompt() {\n content.merge_prompt(|v: &str| self.prompt.replace(INPUT_PLACEHOLDER, v));\n vec![Message::new(MessageRole::User, content)]\n } else {\n let mut messages = vec![];\n let (system, cases) = parse_structure_prompt(&self.prompt);\n if !system.is_empty() {\n messages.push(Message::new(\n MessageRole::System,\n MessageContent::Text(system.to_string()),\n ));\n }\n if !cases.is_empty() {\n messages.extend(cases.into_iter().flat_map(|(i, o)| {\n vec![\n Message::new(MessageRole::User, MessageContent::Text(i.to_string())),\n Message::new(MessageRole::Assistant, MessageContent::Text(o.to_string())),\n ]\n }));\n }\n messages.push(Message::new(MessageRole::User, content));\n messages\n };\n if let Some(text) = input.continue_output() {\n messages.push(Message::new(\n MessageRole::Assistant,\n MessageContent::Text(text.into()),\n ));\n }\n messages\n }\n}\n\nimpl RoleLike for Role {\n fn to_role(&self) -> Role {\n self.clone()\n }\n\n fn model(&self) -> &Model {\n &self.model\n }\n\n fn temperature(&self) -> Option {\n self.temperature\n }\n\n fn top_p(&self) -> Option {\n self.top_p\n }\n\n fn use_tools(&self) -> Option {\n self.use_tools.clone()\n }\n\n fn set_model(&mut self, model: Model) {\n if !self.model().id().is_empty() {\n self.model_id = Some(model.id().to_string());\n }\n self.model = model;\n }\n\n fn set_temperature(&mut self, value: Option) {\n self.temperature = value;\n }\n\n fn set_top_p(&mut self, value: Option) {\n self.top_p = value;\n }\n\n fn set_use_tools(&mut self, value: Option) {\n self.use_tools = value;\n }\n}\n\nfn parse_structure_prompt(prompt: &str) -> (&str, Vec<(&str, &str)>) {\n let mut text = prompt;\n let mut search_input = true;\n let mut system = None;\n let mut parts = vec![];\n loop {\n let search = if search_input {\n \"### INPUT:\"\n } else {\n \"### OUTPUT:\"\n };\n match text.find(search) {\n Some(idx) => {\n if system.is_none() {\n system = Some(&text[..idx])\n } else {\n parts.push(&text[..idx])\n }\n search_input = !search_input;\n text = &text[(idx + search.len())..];\n }\n None => {\n if !text.is_empty() {\n if system.is_none() {\n system = Some(text)\n } else {\n parts.push(text)\n }\n }\n break;\n }\n }\n }\n let parts_len = parts.len();\n if parts_len > 0 && parts_len % 2 == 0 {\n let cases: Vec<(&str, &str)> = parts\n .iter()\n .step_by(2)\n .zip(parts.iter().skip(1).step_by(2))\n .map(|(i, o)| (i.trim(), o.trim()))\n .collect();\n let system = system.map(|v| v.trim()).unwrap_or_default();\n return (system, cases);\n }\n\n (prompt, vec![])\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn test_parse_structure_prompt1() {\n let prompt = r#\"\nSystem message\n### INPUT:\nInput 1\n### OUTPUT:\nOutput 1\n\"#;\n assert_eq!(\n parse_structure_prompt(prompt),\n (\"System message\", vec![(\"Input 1\", \"Output 1\")])\n );\n }\n\n #[test]\n fn test_parse_structure_prompt2() {\n let prompt = r#\"\n### INPUT:\nInput 1\n### OUTPUT:\nOutput 1\n\"#;\n assert_eq!(\n parse_structure_prompt(prompt),\n (\"\", vec![(\"Input 1\", \"Output 1\")])\n );\n }\n\n #[test]\n fn test_parse_structure_prompt3() {\n let prompt = r#\"\nSystem message\n### INPUT:\nInput 1\n\"#;\n assert_eq!(parse_structure_prompt(prompt), (prompt, vec![]));\n }\n}\n" }, { "path": "src/config/session.rs", "content": "use super::input::*;\nuse super::*;\n\nuse crate::client::{Message, MessageContent, MessageRole};\nuse crate::render::MarkdownRender;\n\nuse anyhow::{bail, Context, Result};\nuse fancy_regex::Regex;\nuse inquire::{validator::Validation, Confirm, Text};\nuse serde::{Deserialize, Serialize};\nuse serde_json::json;\nuse std::collections::HashMap;\nuse std::fs::{read_to_string, write};\nuse std::path::Path;\nuse std::sync::LazyLock;\n\nstatic RE_AUTONAME_PREFIX: LazyLock = LazyLock::new(|| Regex::new(r\"\\d{8}T\\d{6}-\").unwrap());\n\n#[derive(Debug, Clone, Default, Deserialize, Serialize)]\npub struct Session {\n #[serde(rename(serialize = \"model\", deserialize = \"model\"))]\n model_id: String,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n temperature: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n top_p: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n use_tools: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n save_session: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n compress_threshold: Option,\n\n #[serde(skip_serializing_if = \"Option::is_none\")]\n role_name: Option,\n #[serde(default, skip_serializing_if = \"IndexMap::is_empty\")]\n agent_variables: AgentVariables,\n #[serde(default, skip_serializing_if = \"String::is_empty\")]\n agent_instructions: String,\n\n #[serde(default, skip_serializing_if = \"Vec::is_empty\")]\n compressed_messages: Vec,\n messages: Vec,\n #[serde(default, skip_serializing_if = \"HashMap::is_empty\")]\n data_urls: HashMap,\n\n #[serde(skip)]\n model: Model,\n #[serde(skip)]\n role_prompt: String,\n #[serde(skip)]\n name: String,\n #[serde(skip)]\n path: Option,\n #[serde(skip)]\n dirty: bool,\n #[serde(skip)]\n save_session_this_time: bool,\n #[serde(skip)]\n compressing: bool,\n #[serde(skip)]\n autoname: Option,\n #[serde(skip)]\n tokens: usize,\n}\n\nimpl Session {\n pub fn new(config: &Config, name: &str) -> Self {\n let role = config.extract_role();\n let mut session = Self {\n name: name.to_string(),\n save_session: config.save_session,\n ..Default::default()\n };\n session.set_role(role);\n session.dirty = false;\n session\n }\n\n pub fn load(config: &Config, name: &str, path: &Path) -> Result {\n let content = read_to_string(path)\n .with_context(|| format!(\"Failed to load session {} at {}\", name, path.display()))?;\n let mut session: Self =\n serde_yaml::from_str(&content).with_context(|| format!(\"Invalid session {name}\"))?;\n\n session.model = Model::retrieve_model(config, &session.model_id, ModelType::Chat)?;\n\n if let Some(autoname) = name.strip_prefix(\"_/\") {\n session.name = TEMP_SESSION_NAME.to_string();\n session.path = None;\n if let Ok(true) = RE_AUTONAME_PREFIX.is_match(autoname) {\n session.autoname = Some(AutoName::new(autoname[16..].to_string()));\n }\n } else {\n session.name = name.to_string();\n session.path = Some(path.display().to_string());\n }\n\n if let Some(role_name) = &session.role_name {\n if let Ok(role) = config.retrieve_role(role_name) {\n session.role_prompt = role.prompt().to_string();\n }\n }\n\n session.update_tokens();\n\n Ok(session)\n }\n\n pub fn is_empty(&self) -> bool {\n self.messages.is_empty() && self.compressed_messages.is_empty()\n }\n\n pub fn name(&self) -> &str {\n &self.name\n }\n\n pub fn role_name(&self) -> Option<&str> {\n self.role_name.as_deref()\n }\n\n pub fn dirty(&self) -> bool {\n self.dirty\n }\n\n pub fn save_session(&self) -> Option {\n self.save_session\n }\n\n pub fn tokens(&self) -> usize {\n self.tokens\n }\n\n pub fn update_tokens(&mut self) {\n self.tokens = self.model().total_tokens(&self.messages);\n }\n\n pub fn has_user_messages(&self) -> bool {\n self.messages.iter().any(|v| v.role.is_user())\n }\n\n pub fn user_messages_len(&self) -> usize {\n self.messages.iter().filter(|v| v.role.is_user()).count()\n }\n\n pub fn export(&self) -> Result {\n let mut data = json!({\n \"path\": self.path,\n \"model\": self.model().id(),\n });\n if let Some(temperature) = self.temperature() {\n data[\"temperature\"] = temperature.into();\n }\n if let Some(top_p) = self.top_p() {\n data[\"top_p\"] = top_p.into();\n }\n if let Some(use_tools) = self.use_tools() {\n data[\"use_tools\"] = use_tools.into();\n }\n if let Some(save_session) = self.save_session() {\n data[\"save_session\"] = save_session.into();\n }\n let (tokens, percent) = self.tokens_usage();\n data[\"total_tokens\"] = tokens.into();\n if let Some(max_input_tokens) = self.model().max_input_tokens() {\n data[\"max_input_tokens\"] = max_input_tokens.into();\n }\n if percent != 0.0 {\n data[\"total/max\"] = format!(\"{percent}%\").into();\n }\n data[\"messages\"] = json!(self.messages);\n\n let output = serde_yaml::to_string(&data)\n .with_context(|| format!(\"Unable to show info about session '{}'\", &self.name))?;\n Ok(output)\n }\n\n pub fn render(\n &self,\n render: &mut MarkdownRender,\n agent_info: &Option<(String, Vec)>,\n ) -> Result {\n let mut items = vec![];\n\n if let Some(path) = &self.path {\n items.push((\"path\", path.to_string()));\n }\n\n if let Some(autoname) = self.autoname() {\n items.push((\"autoname\", autoname.to_string()));\n }\n\n items.push((\"model\", self.model().id()));\n\n if let Some(temperature) = self.temperature() {\n items.push((\"temperature\", temperature.to_string()));\n }\n if let Some(top_p) = self.top_p() {\n items.push((\"top_p\", top_p.to_string()));\n }\n\n if let Some(use_tools) = self.use_tools() {\n items.push((\"use_tools\", use_tools));\n }\n\n if let Some(save_session) = self.save_session() {\n items.push((\"save_session\", save_session.to_string()));\n }\n\n if let Some(compress_threshold) = self.compress_threshold {\n items.push((\"compress_threshold\", compress_threshold.to_string()));\n }\n\n if let Some(max_input_tokens) = self.model().max_input_tokens() {\n items.push((\"max_input_tokens\", max_input_tokens.to_string()));\n }\n\n let mut lines: Vec = items\n .iter()\n .map(|(name, value)| format!(\"{name:<20}{value}\"))\n .collect();\n\n lines.push(String::new());\n\n if !self.is_empty() {\n let resolve_url_fn = |url: &str| resolve_data_url(&self.data_urls, url.to_string());\n\n for message in &self.messages {\n match message.role {\n MessageRole::System => {\n lines.push(\n render\n .render(&message.content.render_input(resolve_url_fn, agent_info)),\n );\n }\n MessageRole::Assistant => {\n if let MessageContent::Text(text) = &message.content {\n lines.push(render.render(text));\n }\n lines.push(\"\".into());\n }\n MessageRole::User => {\n lines.push(format!(\n \">> {}\",\n message.content.render_input(resolve_url_fn, agent_info)\n ));\n }\n MessageRole::Tool => {\n lines.push(message.content.render_input(resolve_url_fn, agent_info));\n }\n }\n }\n }\n\n Ok(lines.join(\"\\n\"))\n }\n\n pub fn tokens_usage(&self) -> (usize, f32) {\n let tokens = self.tokens();\n let max_input_tokens = self.model().max_input_tokens().unwrap_or_default();\n let percent = if max_input_tokens == 0 {\n 0.0\n } else {\n let percent = tokens as f32 / max_input_tokens as f32 * 100.0;\n (percent * 100.0).round() / 100.0\n };\n (tokens, percent)\n }\n\n pub fn set_role(&mut self, role: Role) {\n self.model_id = role.model().id();\n self.temperature = role.temperature();\n self.top_p = role.top_p();\n self.use_tools = role.use_tools();\n self.model = role.model().clone();\n self.role_name = convert_option_string(role.name());\n self.role_prompt = role.prompt().to_string();\n self.dirty = true;\n self.update_tokens();\n }\n\n pub fn clear_role(&mut self) {\n self.role_name = None;\n self.role_prompt.clear();\n }\n\n pub fn sync_agent(&mut self, agent: &Agent) {\n self.role_name = None;\n self.role_prompt = agent.interpolated_instructions();\n self.agent_variables = agent.variables().clone();\n self.agent_instructions = self.role_prompt.clone();\n }\n\n pub fn agent_variables(&self) -> &AgentVariables {\n &self.agent_variables\n }\n\n pub fn agent_instructions(&self) -> &str {\n &self.agent_instructions\n }\n\n pub fn set_save_session(&mut self, value: Option) {\n if self.save_session != value {\n self.save_session = value;\n self.dirty = true;\n }\n }\n\n pub fn set_save_session_this_time(&mut self) {\n self.save_session_this_time = true;\n }\n\n pub fn set_compress_threshold(&mut self, value: Option) {\n if self.compress_threshold != value {\n self.compress_threshold = value;\n self.dirty = true;\n }\n }\n\n pub fn need_compress(&self, global_compress_threshold: usize) -> bool {\n if self.compressing {\n return false;\n }\n let threshold = self.compress_threshold.unwrap_or(global_compress_threshold);\n if threshold < 1 {\n return false;\n }\n self.tokens() > threshold\n }\n\n pub fn compressing(&self) -> bool {\n self.compressing\n }\n\n pub fn set_compressing(&mut self, compressing: bool) {\n self.compressing = compressing;\n }\n\n pub fn compress(&mut self, mut prompt: String) {\n if let Some(system_prompt) = self.messages.first().and_then(|v| {\n if MessageRole::System == v.role {\n let content = v.content.to_text();\n if !content.is_empty() {\n return Some(content);\n }\n }\n None\n }) {\n prompt = format!(\"{system_prompt}\\n\\n{prompt}\",);\n }\n self.compressed_messages.append(&mut self.messages);\n self.messages.push(Message::new(\n MessageRole::System,\n MessageContent::Text(prompt),\n ));\n self.dirty = true;\n self.update_tokens();\n }\n\n pub fn need_autoname(&self) -> bool {\n self.autoname.as_ref().map(|v| v.need()).unwrap_or_default()\n }\n\n pub fn set_autonaming(&mut self, naming: bool) {\n if let Some(v) = self.autoname.as_mut() {\n v.naming = naming;\n }\n }\n\n pub fn chat_history_for_autonaming(&self) -> Option {\n self.autoname.as_ref().and_then(|v| v.chat_history.clone())\n }\n\n pub fn autoname(&self) -> Option<&str> {\n self.autoname.as_ref().and_then(|v| v.name.as_deref())\n }\n\n pub fn set_autoname(&mut self, value: &str) {\n let name = value\n .chars()\n .map(|v| if v.is_alphanumeric() { v } else { '-' })\n .collect();\n self.autoname = Some(AutoName::new(name));\n }\n\n pub fn exit(&mut self, session_dir: &Path, is_repl: bool) -> Result<()> {\n let mut save_session = self.save_session();\n if self.save_session_this_time {\n save_session = Some(true);\n }\n if self.dirty && save_session != Some(false) {\n let mut session_dir = session_dir.to_path_buf();\n let mut session_name = self.name().to_string();\n if save_session.is_none() {\n if !is_repl {\n return Ok(());\n }\n let ans = Confirm::new(\"Save session?\").with_default(false).prompt()?;\n if !ans {\n return Ok(());\n }\n if session_name == TEMP_SESSION_NAME {\n session_name = Text::new(\"Session name:\")\n .with_validator(|input: &str| {\n let input = input.trim();\n if input.is_empty() {\n Ok(Validation::Invalid(\"This name is required\".into()))\n } else if input == TEMP_SESSION_NAME {\n Ok(Validation::Invalid(\"This name is reserved\".into()))\n } else {\n Ok(Validation::Valid)\n }\n })\n .prompt()?;\n }\n } else if save_session == Some(true) && session_name == TEMP_SESSION_NAME {\n session_dir = session_dir.join(\"_\");\n ensure_parent_exists(&session_dir).with_context(|| {\n format!(\"Failed to create directory '{}'\", session_dir.display())\n })?;\n\n let now = chrono::Local::now();\n session_name = now.format(\"%Y%m%dT%H%M%S\").to_string();\n if let Some(autoname) = self.autoname() {\n session_name = format!(\"{session_name}-{autoname}\")\n }\n }\n let session_path = session_dir.join(format!(\"{session_name}.yaml\"));\n self.save(&session_name, &session_path, is_repl)?;\n }\n Ok(())\n }\n\n pub fn save(&mut self, session_name: &str, session_path: &Path, is_repl: bool) -> Result<()> {\n ensure_parent_exists(session_path)?;\n\n self.path = Some(session_path.display().to_string());\n\n let content = serde_yaml::to_string(&self)\n .with_context(|| format!(\"Failed to serde session '{}'\", self.name))?;\n write(session_path, content).with_context(|| {\n format!(\n \"Failed to write session '{}' to '{}'\",\n self.name,\n session_path.display()\n )\n })?;\n\n if is_repl {\n println!(\"✓ Saved the session to '{}'.\", session_path.display());\n }\n\n if self.name() != session_name {\n self.name = session_name.to_string()\n }\n\n self.dirty = false;\n\n Ok(())\n }\n\n pub fn guard_empty(&self) -> Result<()> {\n if !self.is_empty() {\n bail!(\"Cannot perform this operation because the session has messages, please `.empty session` first.\");\n }\n Ok(())\n }\n\n pub fn add_message(&mut self, input: &Input, output: &str) -> Result<()> {\n if input.continue_output().is_some() {\n if let Some(message) = self.messages.last_mut() {\n if let MessageContent::Text(text) = &mut message.content {\n *text = format!(\"{text}{output}\");\n }\n }\n } else if input.regenerate() {\n if let Some(message) = self.messages.last_mut() {\n if let MessageContent::Text(text) = &mut message.content {\n *text = output.to_string();\n }\n }\n } else {\n if self.messages.is_empty() {\n if self.name == TEMP_SESSION_NAME && self.save_session == Some(true) {\n let raw_input = input.raw();\n let chat_history = format!(\"USER: {raw_input}\\nASSISTANT: {output}\\n\");\n self.autoname = Some(AutoName::new_from_chat_history(chat_history));\n }\n self.messages.extend(input.role().build_messages(input));\n } else {\n self.messages\n .push(Message::new(MessageRole::User, input.message_content()));\n }\n self.data_urls.extend(input.data_urls());\n if let Some(tool_calls) = input.tool_calls() {\n self.messages.push(Message::new(\n MessageRole::Tool,\n MessageContent::ToolCalls(tool_calls.clone()),\n ))\n }\n self.messages.push(Message::new(\n MessageRole::Assistant,\n MessageContent::Text(output.to_string()),\n ));\n }\n self.dirty = true;\n self.update_tokens();\n Ok(())\n }\n\n pub fn clear_messages(&mut self) {\n self.messages.clear();\n self.compressed_messages.clear();\n self.data_urls.clear();\n self.autoname = None;\n self.dirty = true;\n self.update_tokens();\n }\n\n pub fn echo_messages(&self, input: &Input) -> String {\n let messages = self.build_messages(input);\n serde_yaml::to_string(&messages).unwrap_or_else(|_| \"Unable to echo message\".into())\n }\n\n pub fn build_messages(&self, input: &Input) -> Vec {\n let mut messages = self.messages.clone();\n if input.continue_output().is_some() {\n return messages;\n } else if input.regenerate() {\n while let Some(last) = messages.last() {\n if !last.role.is_user() {\n messages.pop();\n } else {\n break;\n }\n }\n return messages;\n }\n let mut need_add_msg = true;\n let len = messages.len();\n if len == 0 {\n messages = input.role().build_messages(input);\n need_add_msg = false;\n } else if len == 1 && self.compressed_messages.len() >= 2 {\n if let Some(index) = self\n .compressed_messages\n .iter()\n .rposition(|v| v.role == MessageRole::User)\n {\n messages.extend(self.compressed_messages[index..].to_vec());\n }\n }\n if need_add_msg {\n messages.push(Message::new(MessageRole::User, input.message_content()));\n }\n messages\n }\n}\n\nimpl RoleLike for Session {\n fn to_role(&self) -> Role {\n let role_name = self.role_name.as_deref().unwrap_or_default();\n let mut role = Role::new(role_name, &self.role_prompt);\n role.sync(self);\n role\n }\n\n fn model(&self) -> &Model {\n &self.model\n }\n\n fn temperature(&self) -> Option {\n self.temperature\n }\n\n fn top_p(&self) -> Option {\n self.top_p\n }\n\n fn use_tools(&self) -> Option {\n self.use_tools.clone()\n }\n\n fn set_model(&mut self, model: Model) {\n if self.model().id() != model.id() {\n self.model_id = model.id();\n self.model = model;\n self.dirty = true;\n self.update_tokens();\n }\n }\n\n fn set_temperature(&mut self, value: Option) {\n if self.temperature != value {\n self.temperature = value;\n self.dirty = true;\n }\n }\n\n fn set_top_p(&mut self, value: Option) {\n if self.top_p != value {\n self.top_p = value;\n self.dirty = true;\n }\n }\n\n fn set_use_tools(&mut self, value: Option) {\n if self.use_tools != value {\n self.use_tools = value;\n self.dirty = true;\n }\n }\n}\n\n#[derive(Debug, Clone, Default)]\nstruct AutoName {\n naming: bool,\n chat_history: Option,\n name: Option,\n}\n\nimpl AutoName {\n pub fn new(name: String) -> Self {\n Self {\n name: Some(name),\n ..Default::default()\n }\n }\n pub fn new_from_chat_history(chat_history: String) -> Self {\n Self {\n chat_history: Some(chat_history),\n ..Default::default()\n }\n }\n pub fn need(&self) -> bool {\n !self.naming && self.chat_history.is_some() && self.name.is_none()\n }\n}\n" }, { "path": "src/function.rs", "content": "use crate::{\n config::{Agent, Config, GlobalConfig},\n utils::*,\n};\n\nuse anyhow::{anyhow, bail, Context, Result};\nuse indexmap::IndexMap;\nuse serde::{Deserialize, Serialize};\nuse serde_json::{json, Value};\nuse std::{\n collections::{HashMap, HashSet},\n fs,\n path::{Path, PathBuf},\n};\n\n#[cfg(windows)]\nconst PATH_SEP: &str = \";\";\n#[cfg(not(windows))]\nconst PATH_SEP: &str = \":\";\n\npub fn eval_tool_calls(config: &GlobalConfig, mut calls: Vec) -> Result> {\n let mut output = vec![];\n if calls.is_empty() {\n return Ok(output);\n }\n calls = ToolCall::dedup(calls);\n if calls.is_empty() {\n bail!(\"The request was aborted because an infinite loop of function calls was detected.\")\n }\n let mut is_all_null = true;\n for call in calls {\n let mut result = call.eval(config)?;\n if result.is_null() {\n result = json!(\"DONE\");\n } else {\n is_all_null = false;\n }\n output.push(ToolResult::new(call, result));\n }\n if is_all_null {\n output = vec![];\n }\n Ok(output)\n}\n\n#[derive(Debug, Clone, Deserialize, Serialize)]\npub struct ToolResult {\n pub call: ToolCall,\n pub output: Value,\n}\n\nimpl ToolResult {\n pub fn new(call: ToolCall, output: Value) -> Self {\n Self { call, output }\n }\n}\n\n#[derive(Debug, Clone, Default)]\npub struct Functions {\n declarations: Vec,\n}\n\nimpl Functions {\n pub fn init(declarations_path: &Path) -> Result {\n let declarations: Vec = if declarations_path.exists() {\n let ctx = || {\n format!(\n \"Failed to load functions at {}\",\n declarations_path.display()\n )\n };\n let content = fs::read_to_string(declarations_path).with_context(ctx)?;\n serde_json::from_str(&content).with_context(ctx)?\n } else {\n vec![]\n };\n\n Ok(Self { declarations })\n }\n\n pub fn find(&self, name: &str) -> Option<&FunctionDeclaration> {\n self.declarations.iter().find(|v| v.name == name)\n }\n\n pub fn contains(&self, name: &str) -> bool {\n self.declarations.iter().any(|v| v.name == name)\n }\n\n pub fn declarations(&self) -> &[FunctionDeclaration] {\n &self.declarations\n }\n\n pub fn is_empty(&self) -> bool {\n self.declarations.is_empty()\n }\n}\n\n#[derive(Debug, Clone, Serialize, Deserialize)]\npub struct FunctionDeclaration {\n pub name: String,\n pub description: String,\n pub parameters: JsonSchema,\n #[serde(skip_serializing, default)]\n pub agent: bool,\n}\n\n#[derive(Debug, Clone, Serialize, Deserialize)]\npub struct JsonSchema {\n #[serde(rename = \"type\", skip_serializing_if = \"Option::is_none\")]\n pub type_value: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub description: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub properties: Option>,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub items: Option>,\n #[serde(rename = \"anyOf\", skip_serializing_if = \"Option::is_none\")]\n pub any_of: Option>,\n #[serde(rename = \"enum\", skip_serializing_if = \"Option::is_none\")]\n pub enum_value: Option>,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub default: Option,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub required: Option>,\n}\n\nimpl JsonSchema {\n pub fn is_empty_properties(&self) -> bool {\n match &self.properties {\n Some(v) => v.is_empty(),\n None => true,\n }\n }\n}\n\n#[derive(Debug, Clone, Default, Deserialize, Serialize)]\npub struct ToolCall {\n pub name: String,\n pub arguments: Value,\n pub id: Option,\n}\n\ntype CallConfig = (String, String, Vec, HashMap);\n\nimpl ToolCall {\n pub fn dedup(calls: Vec) -> Vec {\n let mut new_calls = vec![];\n let mut seen_ids = HashSet::new();\n\n for call in calls.into_iter().rev() {\n if let Some(id) = &call.id {\n if !seen_ids.contains(id) {\n seen_ids.insert(id.clone());\n new_calls.push(call);\n }\n } else {\n new_calls.push(call);\n }\n }\n\n new_calls.reverse();\n new_calls\n }\n\n pub fn new(name: String, arguments: Value, id: Option) -> Self {\n Self {\n name,\n arguments,\n id,\n }\n }\n\n pub fn eval(&self, config: &GlobalConfig) -> Result {\n let (call_name, cmd_name, mut cmd_args, envs) = match &config.read().agent {\n Some(agent) => self.extract_call_config_from_agent(config, agent)?,\n None => self.extract_call_config_from_config(config)?,\n };\n\n let json_data = if self.arguments.is_object() {\n self.arguments.clone()\n } else if let Some(arguments) = self.arguments.as_str() {\n let arguments: Value = serde_json::from_str(arguments).map_err(|_| {\n anyhow!(\"The call '{call_name}' has invalid arguments: {arguments}\")\n })?;\n arguments\n } else {\n bail!(\n \"The call '{call_name}' has invalid arguments: {}\",\n self.arguments\n );\n };\n\n cmd_args.push(json_data.to_string());\n\n let output = match run_llm_function(cmd_name, cmd_args, envs)? {\n Some(contents) => serde_json::from_str(&contents)\n .ok()\n .unwrap_or_else(|| json!({\"output\": contents})),\n None => Value::Null,\n };\n\n Ok(output)\n }\n\n fn extract_call_config_from_agent(\n &self,\n config: &GlobalConfig,\n agent: &Agent,\n ) -> Result {\n let function_name = self.name.clone();\n match agent.functions().find(&function_name) {\n Some(function) => {\n let agent_name = agent.name().to_string();\n if function.agent {\n Ok((\n format!(\"{agent_name}-{function_name}\"),\n agent_name,\n vec![function_name],\n agent.variable_envs(),\n ))\n } else {\n Ok((\n function_name.clone(),\n function_name,\n vec![],\n Default::default(),\n ))\n }\n }\n None => self.extract_call_config_from_config(config),\n }\n }\n\n fn extract_call_config_from_config(&self, config: &GlobalConfig) -> Result {\n let function_name = self.name.clone();\n match config.read().functions.contains(&function_name) {\n true => Ok((\n function_name.clone(),\n function_name,\n vec![],\n Default::default(),\n )),\n false => bail!(\"Unexpected call: {function_name} {}\", self.arguments),\n }\n }\n}\n\npub fn run_llm_function(\n cmd_name: String,\n cmd_args: Vec,\n mut envs: HashMap,\n) -> Result> {\n let prompt = format!(\"Call {cmd_name} {}\", cmd_args.join(\" \"));\n\n let mut bin_dirs: Vec = vec![];\n if cmd_args.len() > 1 {\n let dir = Config::agent_functions_dir(&cmd_name).join(\"bin\");\n if dir.exists() {\n bin_dirs.push(dir);\n }\n }\n bin_dirs.push(Config::functions_bin_dir());\n let current_path = std::env::var(\"PATH\").context(\"No PATH environment variable\")?;\n let prepend_path = bin_dirs\n .iter()\n .map(|v| format!(\"{}{PATH_SEP}\", v.display()))\n .collect::>()\n .join(\"\");\n envs.insert(\"PATH\".into(), format!(\"{prepend_path}{current_path}\"));\n\n let temp_file = temp_file(\"-eval-\", \"\");\n envs.insert(\"LLM_OUTPUT\".into(), temp_file.display().to_string());\n\n #[cfg(windows)]\n let cmd_name = polyfill_cmd_name(&cmd_name, &bin_dirs);\n if *IS_STDOUT_TERMINAL {\n println!(\"{}\", dimmed_text(&prompt));\n }\n let exit_code = run_command(&cmd_name, &cmd_args, Some(envs))\n .map_err(|err| anyhow!(\"Unable to run {cmd_name}, {err}\"))?;\n if exit_code != 0 {\n bail!(\"Tool call exit with {exit_code}\");\n }\n let mut output = None;\n if temp_file.exists() {\n let contents =\n fs::read_to_string(temp_file).context(\"Failed to retrieve tool call output\")?;\n if !contents.is_empty() {\n output = Some(contents);\n }\n };\n Ok(output)\n}\n\n#[cfg(windows)]\nfn polyfill_cmd_name>(cmd_name: &str, bin_dir: &[T]) -> String {\n let cmd_name = cmd_name.to_string();\n if let Ok(exts) = std::env::var(\"PATHEXT\") {\n for name in exts.split(';').map(|ext| format!(\"{cmd_name}{ext}\")) {\n for dir in bin_dir {\n let path = dir.as_ref().join(&name);\n if path.exists() {\n return name.to_string();\n }\n }\n }\n }\n cmd_name\n}\n" }, { "path": "src/main.rs", "content": "mod cli;\nmod client;\nmod config;\nmod function;\nmod rag;\nmod render;\nmod repl;\nmod serve;\n#[macro_use]\nmod utils;\n\n#[macro_use]\nextern crate log;\n\nuse crate::cli::Cli;\nuse crate::client::{\n call_chat_completions, call_chat_completions_streaming, list_models, ModelType,\n};\nuse crate::config::{\n ensure_parent_exists, list_agents, load_env_file, macro_execute, Config, GlobalConfig, Input,\n WorkingMode, CODE_ROLE, EXPLAIN_SHELL_ROLE, SHELL_ROLE, TEMP_SESSION_NAME,\n};\nuse crate::render::render_error;\nuse crate::repl::Repl;\nuse crate::utils::*;\n\nuse anyhow::{bail, Result};\nuse clap::Parser;\nuse inquire::Text;\nuse parking_lot::RwLock;\nuse simplelog::{format_description, ConfigBuilder, LevelFilter, SimpleLogger, WriteLogger};\nuse std::{env, process, sync::Arc};\n\n#[tokio::main]\nasync fn main() -> Result<()> {\n load_env_file()?;\n let cli = Cli::parse();\n let text = cli.text()?;\n let working_mode = if cli.serve.is_some() {\n WorkingMode::Serve\n } else if text.is_none() && cli.file.is_empty() {\n WorkingMode::Repl\n } else {\n WorkingMode::Cmd\n };\n let info_flag = cli.info\n || cli.sync_models\n || cli.list_models\n || cli.list_roles\n || cli.list_agents\n || cli.list_rags\n || cli.list_macros\n || cli.list_sessions;\n setup_logger(working_mode.is_serve())?;\n let config = Arc::new(RwLock::new(Config::init(working_mode, info_flag).await?));\n if let Err(err) = run(config, cli, text).await {\n render_error(err);\n std::process::exit(1);\n }\n Ok(())\n}\n\nasync fn run(config: GlobalConfig, cli: Cli, text: Option) -> Result<()> {\n let abort_signal = create_abort_signal();\n\n if cli.sync_models {\n let url = config.read().sync_models_url();\n return Config::sync_models(&url, abort_signal.clone()).await;\n }\n\n if cli.list_models {\n for model in list_models(&config.read(), ModelType::Chat) {\n println!(\"{}\", model.id());\n }\n return Ok(());\n }\n if cli.list_roles {\n let roles = Config::list_roles(true).join(\"\\n\");\n println!(\"{roles}\");\n return Ok(());\n }\n if cli.list_agents {\n let agents = list_agents().join(\"\\n\");\n println!(\"{agents}\");\n return Ok(());\n }\n if cli.list_rags {\n let rags = Config::list_rags().join(\"\\n\");\n println!(\"{rags}\");\n return Ok(());\n }\n if cli.list_macros {\n let macros = Config::list_macros().join(\"\\n\");\n println!(\"{macros}\");\n return Ok(());\n }\n\n if cli.dry_run {\n config.write().dry_run = true;\n }\n\n if let Some(agent) = &cli.agent {\n let session = cli.session.as_ref().map(|v| match v {\n Some(v) => v.as_str(),\n None => TEMP_SESSION_NAME,\n });\n if !cli.agent_variable.is_empty() {\n config.write().agent_variables = Some(\n cli.agent_variable\n .chunks(2)\n .map(|v| (v[0].to_string(), v[1].to_string()))\n .collect(),\n );\n }\n\n let ret = Config::use_agent(&config, agent, session, abort_signal.clone()).await;\n config.write().agent_variables = None;\n ret?;\n } else {\n if let Some(prompt) = &cli.prompt {\n config.write().use_prompt(prompt)?;\n } else if let Some(name) = &cli.role {\n config.write().use_role(name)?;\n } else if cli.execute {\n config.write().use_role(SHELL_ROLE)?;\n } else if cli.code {\n config.write().use_role(CODE_ROLE)?;\n }\n if let Some(session) = &cli.session {\n config\n .write()\n .use_session(session.as_ref().map(|v| v.as_str()))?;\n }\n if let Some(rag) = &cli.rag {\n Config::use_rag(&config, Some(rag), abort_signal.clone()).await?;\n }\n }\n if cli.list_sessions {\n let sessions = config.read().list_sessions().join(\"\\n\");\n println!(\"{sessions}\");\n return Ok(());\n }\n if let Some(model_id) = &cli.model {\n config.write().set_model(model_id)?;\n }\n if cli.no_stream {\n config.write().stream = false;\n }\n if cli.empty_session {\n config.write().empty_session()?;\n }\n if cli.save_session {\n config.write().set_save_session_this_time()?;\n }\n if cli.info {\n let info = config.read().info()?;\n println!(\"{info}\");\n return Ok(());\n }\n if let Some(addr) = cli.serve {\n return serve::run(config, addr).await;\n }\n let is_repl = config.read().working_mode.is_repl();\n if cli.rebuild_rag {\n Config::rebuild_rag(&config, abort_signal.clone()).await?;\n if is_repl {\n return Ok(());\n }\n }\n if let Some(name) = &cli.macro_name {\n macro_execute(&config, name, text.as_deref(), abort_signal.clone()).await?;\n return Ok(());\n }\n if cli.execute && !is_repl {\n let input = create_input(&config, text, &cli.file, abort_signal.clone()).await?;\n shell_execute(&config, &SHELL, input, abort_signal.clone()).await?;\n return Ok(());\n }\n config.write().apply_prelude()?;\n match is_repl {\n false => {\n let mut input = create_input(&config, text, &cli.file, abort_signal.clone()).await?;\n input.use_embeddings(abort_signal.clone()).await?;\n start_directive(&config, input, cli.code, abort_signal).await\n }\n true => {\n if !*IS_STDOUT_TERMINAL {\n bail!(\"No TTY for REPL\")\n }\n start_interactive(&config).await\n }\n }\n}\n\n#[async_recursion::async_recursion]\nasync fn start_directive(\n config: &GlobalConfig,\n input: Input,\n code_mode: bool,\n abort_signal: AbortSignal,\n) -> Result<()> {\n let client = input.create_client()?;\n let extract_code = !*IS_STDOUT_TERMINAL && code_mode;\n config.write().before_chat_completion(&input)?;\n let (output, tool_results) = if !input.stream() || extract_code {\n call_chat_completions(\n &input,\n true,\n extract_code,\n client.as_ref(),\n abort_signal.clone(),\n )\n .await?\n } else {\n call_chat_completions_streaming(&input, client.as_ref(), abort_signal.clone()).await?\n };\n config\n .write()\n .after_chat_completion(&input, &output, &tool_results)?;\n\n if !tool_results.is_empty() {\n start_directive(\n config,\n input.merge_tool_results(output, tool_results),\n code_mode,\n abort_signal,\n )\n .await?;\n }\n\n config.write().exit_session()?;\n Ok(())\n}\n\nasync fn start_interactive(config: &GlobalConfig) -> Result<()> {\n let mut repl: Repl = Repl::init(config)?;\n repl.run().await\n}\n\n#[async_recursion::async_recursion]\nasync fn shell_execute(\n config: &GlobalConfig,\n shell: &Shell,\n mut input: Input,\n abort_signal: AbortSignal,\n) -> Result<()> {\n let client = input.create_client()?;\n config.write().before_chat_completion(&input)?;\n let (eval_str, _) =\n call_chat_completions(&input, false, true, client.as_ref(), abort_signal.clone()).await?;\n\n config\n .write()\n .after_chat_completion(&input, &eval_str, &[])?;\n if eval_str.is_empty() {\n bail!(\"No command generated\");\n }\n if config.read().dry_run {\n config.read().print_markdown(&eval_str)?;\n return Ok(());\n }\n if *IS_STDOUT_TERMINAL {\n let options = [\"execute\", \"revise\", \"describe\", \"copy\", \"quit\"];\n let command = color_text(eval_str.trim(), nu_ansi_term::Color::Rgb(255, 165, 0));\n let first_letter_color = nu_ansi_term::Color::Cyan;\n let prompt_text = options\n .iter()\n .map(|v| format!(\"{}{}\", color_text(&v[0..1], first_letter_color), &v[1..]))\n .collect::>()\n .join(&dimmed_text(\" | \"));\n loop {\n println!(\"{command}\");\n let answer_char =\n read_single_key(&['e', 'r', 'd', 'c', 'q'], 'e', &format!(\"{prompt_text}: \"))?;\n\n match answer_char {\n 'e' => {\n debug!(\"{} {:?}\", shell.cmd, &[&shell.arg, &eval_str]);\n let code = run_command(&shell.cmd, &[&shell.arg, &eval_str], None)?;\n if code == 0 && config.read().save_shell_history {\n let _ = append_to_shell_history(&shell.name, &eval_str, code);\n }\n process::exit(code);\n }\n 'r' => {\n let revision = Text::new(\"Enter your revision:\").prompt()?;\n let text = format!(\"{}\\n{revision}\", input.text());\n input.set_text(text);\n return shell_execute(config, shell, input, abort_signal.clone()).await;\n }\n 'd' => {\n let role = config.read().retrieve_role(EXPLAIN_SHELL_ROLE)?;\n let input = Input::from_str(config, &eval_str, Some(role));\n if input.stream() {\n call_chat_completions_streaming(\n &input,\n client.as_ref(),\n abort_signal.clone(),\n )\n .await?;\n } else {\n call_chat_completions(\n &input,\n true,\n false,\n client.as_ref(),\n abort_signal.clone(),\n )\n .await?;\n }\n println!();\n continue;\n }\n 'c' => {\n set_text(&eval_str)?;\n println!(\"{}\", dimmed_text(\"✓ Copied the command.\"));\n }\n _ => {}\n }\n break;\n }\n } else {\n println!(\"{eval_str}\");\n }\n Ok(())\n}\n\nasync fn create_input(\n config: &GlobalConfig,\n text: Option,\n file: &[String],\n abort_signal: AbortSignal,\n) -> Result {\n let input = if file.is_empty() {\n Input::from_str(config, &text.unwrap_or_default(), None)\n } else {\n Input::from_files_with_spinner(\n config,\n &text.unwrap_or_default(),\n file.to_vec(),\n None,\n abort_signal,\n )\n .await?\n };\n if input.is_empty() {\n bail!(\"No input\");\n }\n Ok(input)\n}\n\nfn setup_logger(is_serve: bool) -> Result<()> {\n let (log_level, log_path) = Config::log_config(is_serve)?;\n if log_level == LevelFilter::Off {\n return Ok(());\n }\n let crate_name = env!(\"CARGO_CRATE_NAME\");\n let log_filter = match std::env::var(get_env_name(\"log_filter\")) {\n Ok(v) => v,\n Err(_) => match is_serve {\n true => format!(\"{crate_name}::serve\"),\n false => crate_name.into(),\n },\n };\n let config = ConfigBuilder::new()\n .add_filter_allow(log_filter)\n .set_time_format_custom(format_description!(\n \"[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z\"\n ))\n .set_thread_level(LevelFilter::Off)\n .build();\n match log_path {\n None => {\n SimpleLogger::init(log_level, config)?;\n }\n Some(log_path) => {\n ensure_parent_exists(&log_path)?;\n let log_file = std::fs::File::create(log_path)?;\n WriteLogger::init(log_level, config, log_file)?;\n }\n }\n Ok(())\n}\n" }, { "path": "src/rag/mod.rs", "content": "use self::splitter::*;\n\nuse crate::client::*;\nuse crate::config::*;\nuse crate::utils::*;\n\nmod serde_vectors;\nmod splitter;\n\nuse anyhow::{anyhow, bail, Context, Result};\nuse bm25::{Language, SearchEngine, SearchEngineBuilder};\nuse hnsw_rs::prelude::*;\nuse indexmap::{IndexMap, IndexSet};\nuse inquire::{required, validator::Validation, Confirm, Select, Text};\nuse parking_lot::RwLock;\nuse serde::{Deserialize, Serialize};\nuse serde_json::json;\nuse std::{collections::HashMap, env, fmt::Debug, fs, hash::Hash, path::Path, time::Duration};\nuse tokio::time::sleep;\n\npub struct Rag {\n config: GlobalConfig,\n name: String,\n path: String,\n embedding_model: Model,\n hnsw: Hnsw<'static, f32, DistCosine>,\n bm25: SearchEngine,\n data: RagData,\n last_sources: RwLock>,\n}\n\nimpl Debug for Rag {\n fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {\n f.debug_struct(\"Rag\")\n .field(\"name\", &self.name)\n .field(\"path\", &self.path)\n .field(\"embedding_model\", &self.embedding_model)\n .field(\"data\", &self.data)\n .finish()\n }\n}\n\nimpl Clone for Rag {\n fn clone(&self) -> Self {\n Self {\n config: self.config.clone(),\n name: self.name.clone(),\n path: self.path.clone(),\n embedding_model: self.embedding_model.clone(),\n hnsw: self.data.build_hnsw(),\n bm25: self.data.build_bm25(),\n data: self.data.clone(),\n last_sources: RwLock::new(None),\n }\n }\n}\n\nimpl Rag {\n pub async fn init(\n config: &GlobalConfig,\n name: &str,\n save_path: &Path,\n doc_paths: &[String],\n abort_signal: AbortSignal,\n ) -> Result {\n if !*IS_STDOUT_TERMINAL {\n bail!(\"Failed to init rag in non-interactive mode\");\n }\n println!(\"⚙ Initializing RAG...\");\n let (embedding_model, chunk_size, chunk_overlap) = Self::create_config(config)?;\n let (reranker_model, top_k) = {\n let config = config.read();\n (config.rag_reranker_model.clone(), config.rag_top_k)\n };\n let data = RagData::new(\n embedding_model.id(),\n chunk_size,\n chunk_overlap,\n reranker_model,\n top_k,\n embedding_model.max_batch_size(),\n );\n let mut rag = Self::create(config, name, save_path, data)?;\n let mut paths = doc_paths.to_vec();\n if paths.is_empty() {\n paths = add_documents()?;\n };\n let loaders = config.read().document_loaders.clone();\n let (spinner, spinner_rx) = Spinner::create(\"\");\n abortable_run_with_spinner_rx(\n rag.sync_documents(&paths, true, loaders, Some(spinner)),\n spinner_rx,\n abort_signal,\n )\n .await?;\n if rag.save()? {\n println!(\"✓ Saved RAG to '{}'.\", save_path.display());\n }\n Ok(rag)\n }\n\n pub fn load(config: &GlobalConfig, name: &str, path: &Path) -> Result {\n let err = || format!(\"Failed to load rag '{name}' at '{}'\", path.display());\n let content = fs::read_to_string(path).with_context(err)?;\n let data: RagData = serde_yaml::from_str(&content).with_context(err)?;\n Self::create(config, name, path, data)\n }\n\n pub fn create(config: &GlobalConfig, name: &str, path: &Path, data: RagData) -> Result {\n let hnsw = data.build_hnsw();\n let bm25 = data.build_bm25();\n let embedding_model =\n Model::retrieve_model(&config.read(), &data.embedding_model, ModelType::Embedding)?;\n let rag = Rag {\n config: config.clone(),\n name: name.to_string(),\n path: path.display().to_string(),\n data,\n embedding_model,\n hnsw,\n bm25,\n last_sources: RwLock::new(None),\n };\n Ok(rag)\n }\n\n pub fn document_paths(&self) -> &[String] {\n &self.data.document_paths\n }\n\n pub async fn refresh_document_paths(\n &mut self,\n document_paths: &[String],\n refresh: bool,\n config: &GlobalConfig,\n abort_signal: AbortSignal,\n ) -> Result<()> {\n let loaders = config.read().document_loaders.clone();\n let (spinner, spinner_rx) = Spinner::create(\"\");\n abortable_run_with_spinner_rx(\n self.sync_documents(document_paths, refresh, loaders, Some(spinner)),\n spinner_rx,\n abort_signal,\n )\n .await?;\n if self.save()? {\n println!(\"✓ Saved rag to '{}'.\", self.path);\n }\n Ok(())\n }\n\n pub fn create_config(config: &GlobalConfig) -> Result<(Model, usize, usize)> {\n let (embedding_model_id, chunk_size, chunk_overlap) = {\n let config = config.read();\n (\n config.rag_embedding_model.clone(),\n config.rag_chunk_size,\n config.rag_chunk_overlap,\n )\n };\n let embedding_model_id = match embedding_model_id {\n Some(value) => {\n println!(\"Select embedding model: {value}\");\n value\n }\n None => {\n let models = list_models(&config.read(), ModelType::Embedding);\n if models.is_empty() {\n bail!(\"No available embedding model\");\n }\n select_embedding_model(&models)?\n }\n };\n let embedding_model =\n Model::retrieve_model(&config.read(), &embedding_model_id, ModelType::Embedding)?;\n\n let chunk_size = match chunk_size {\n Some(value) => {\n println!(\"Set chunk size: {value}\");\n value\n }\n None => set_chunk_size(&embedding_model)?,\n };\n let chunk_overlap = match chunk_overlap {\n Some(value) => {\n println!(\"Set chunk overlay: {value}\");\n value\n }\n None => {\n let value = chunk_size / 20;\n set_chunk_overlay(value)?\n }\n };\n\n Ok((embedding_model, chunk_size, chunk_overlap))\n }\n\n pub fn get_config(&self) -> (Option, usize) {\n (self.data.reranker_model.clone(), self.data.top_k)\n }\n\n pub fn get_last_sources(&self) -> Option {\n self.last_sources.read().clone()\n }\n\n pub fn set_last_sources(&self, ids: &[DocumentId]) {\n let mut sources: IndexMap> = IndexMap::new();\n for id in ids {\n let (file_index, _) = id.split();\n if let Some(file) = self.data.files.get(&file_index) {\n sources\n .entry(file.path.clone())\n .or_default()\n .push(format!(\"{id:?}\"));\n }\n }\n let sources = if sources.is_empty() {\n None\n } else {\n Some(\n sources\n .into_iter()\n .map(|(path, ids)| format!(\"{path} ({})\", ids.join(\",\")))\n .collect::>()\n .join(\"\\n\"),\n )\n };\n *self.last_sources.write() = sources;\n }\n\n pub fn set_reranker_model(&mut self, reranker_model: Option) -> Result<()> {\n self.data.reranker_model = reranker_model;\n self.save()?;\n Ok(())\n }\n\n pub fn set_top_k(&mut self, top_k: usize) -> Result<()> {\n self.data.top_k = top_k;\n self.save()?;\n Ok(())\n }\n\n pub fn save(&self) -> Result {\n if self.is_temp() {\n return Ok(false);\n }\n let path = Path::new(&self.path);\n ensure_parent_exists(path)?;\n\n let content = serde_yaml::to_string(&self.data)\n .with_context(|| format!(\"Failed to serde rag '{}'\", self.name))?;\n fs::write(path, content).with_context(|| {\n format!(\"Failed to save rag '{}' to '{}'\", self.name, path.display())\n })?;\n\n Ok(true)\n }\n\n pub fn export(&self) -> Result {\n let files: Vec<_> = self\n .data\n .files\n .iter()\n .map(|(_, v)| {\n json!({\n \"path\": v.path,\n \"num_chunks\": v.documents.len(),\n })\n })\n .collect();\n let data = json!({\n \"path\": self.path,\n \"embedding_model\": self.embedding_model.id(),\n \"chunk_size\": self.data.chunk_size,\n \"chunk_overlap\": self.data.chunk_overlap,\n \"reranker_model\": self.data.reranker_model,\n \"top_k\": self.data.top_k,\n \"batch_size\": self.data.batch_size,\n \"document_paths\": self.data.document_paths,\n \"files\": files,\n });\n let output = serde_yaml::to_string(&data)\n .with_context(|| format!(\"Unable to show info about rag '{}'\", self.name))?;\n Ok(output)\n }\n\n pub fn name(&self) -> &str {\n &self.name\n }\n\n pub fn is_temp(&self) -> bool {\n self.name == TEMP_RAG_NAME\n }\n\n pub async fn search(\n &self,\n text: &str,\n top_k: usize,\n rerank_model: Option<&str>,\n abort_signal: AbortSignal,\n ) -> Result<(String, Vec)> {\n let ret = abortable_run_with_spinner(\n self.hybird_search(text, top_k, rerank_model),\n \"Searching\",\n abort_signal,\n )\n .await;\n let (ids, documents): (Vec<_>, Vec<_>) = ret?.into_iter().unzip();\n let embeddings = documents.join(\"\\n\\n\");\n Ok((embeddings, ids))\n }\n\n pub async fn sync_documents(\n &mut self,\n paths: &[String],\n refresh: bool,\n loaders: HashMap,\n spinner: Option,\n ) -> Result<()> {\n if let Some(spinner) = &spinner {\n let _ = spinner.set_message(String::new());\n }\n let (document_paths, mut recursive_urls, mut urls, mut protocol_paths, mut local_paths) =\n resolve_paths(&loaders, paths).await?;\n let mut to_deleted: IndexMap> = Default::default();\n if refresh {\n for (file_id, file) in &self.data.files {\n to_deleted\n .entry(file.hash.clone())\n .or_default()\n .push(*file_id);\n }\n } else {\n let recursive_urls_cloned = recursive_urls.clone();\n let match_recursive_url = |v: &str| {\n recursive_urls_cloned\n .iter()\n .any(|start_url| v.starts_with(start_url))\n };\n recursive_urls = recursive_urls\n .into_iter()\n .filter(|v| !self.data.document_paths.contains(&format!(\"{v}**\")))\n .collect();\n let protocol_paths_cloned = protocol_paths.clone();\n let match_protocol_path =\n |v: &str| protocol_paths_cloned.iter().any(|root| v.starts_with(root));\n protocol_paths = protocol_paths\n .into_iter()\n .filter(|v| !self.data.document_paths.contains(v))\n .collect();\n for (file_id, file) in &self.data.files {\n if is_url(&file.path) {\n if !urls.swap_remove(&file.path) && !match_recursive_url(&file.path) {\n to_deleted\n .entry(file.hash.clone())\n .or_default()\n .push(*file_id);\n }\n } else if is_loader_protocol(&loaders, &file.path) {\n if !match_protocol_path(&file.path) {\n to_deleted\n .entry(file.hash.clone())\n .or_default()\n .push(*file_id);\n }\n } else if !local_paths.swap_remove(&file.path) {\n to_deleted\n .entry(file.hash.clone())\n .or_default()\n .push(*file_id);\n }\n }\n }\n\n let mut loaded_documents = vec![];\n let mut has_error = false;\n let mut index = 0;\n let total = recursive_urls.len() + urls.len() + protocol_paths.len() + local_paths.len();\n let handle_error = |error: anyhow::Error, has_error: &mut bool| {\n println!(\"{}\", warning_text(&format!(\"⚠️ {error}\")));\n *has_error = true;\n };\n for start_url in recursive_urls {\n index += 1;\n println!(\"Load {start_url}** [{index}/{total}]\");\n match load_recursive_url(&loaders, &start_url).await {\n Ok(v) => loaded_documents.extend(v),\n Err(err) => handle_error(err, &mut has_error),\n }\n }\n for url in urls {\n index += 1;\n println!(\"Load {url} [{index}/{total}]\");\n match load_url(&loaders, &url).await {\n Ok(v) => loaded_documents.push(v),\n Err(err) => handle_error(err, &mut has_error),\n }\n }\n for protocol_path in protocol_paths {\n index += 1;\n println!(\"Load {protocol_path} [{index}/{total}]\");\n match load_protocol_path(&loaders, &protocol_path) {\n Ok(v) => loaded_documents.extend(v),\n Err(err) => handle_error(err, &mut has_error),\n }\n }\n for local_path in local_paths {\n index += 1;\n println!(\"Load {local_path} [{index}/{total}]\");\n match load_file(&loaders, &local_path).await {\n Ok(v) => loaded_documents.push(v),\n Err(err) => handle_error(err, &mut has_error),\n }\n }\n\n if has_error {\n let mut aborted = true;\n if *IS_STDOUT_TERMINAL && total > 0 {\n let ans = Confirm::new(\"Some documents failed to load. Continue?\")\n .with_default(false)\n .prompt()?;\n aborted = !ans;\n }\n if aborted {\n bail!(\"Aborted\");\n }\n }\n\n let mut rag_files = vec![];\n for LoadedDocument {\n path,\n contents,\n mut metadata,\n } in loaded_documents\n {\n let hash = sha256(&contents);\n if let Some(file_ids) = to_deleted.get_mut(&hash) {\n if let Some((i, _)) = file_ids\n .iter()\n .enumerate()\n .find(|(_, v)| self.data.files[*v].path == path)\n {\n if file_ids.len() == 1 {\n to_deleted.swap_remove(&hash);\n } else {\n file_ids.remove(i);\n }\n continue;\n }\n }\n let extension = metadata\n .swap_remove(EXTENSION_METADATA)\n .unwrap_or_else(|| DEFAULT_EXTENSION.into());\n let separator = get_separators(&extension);\n let splitter = RecursiveCharacterTextSplitter::new(\n self.data.chunk_size,\n self.data.chunk_overlap,\n &separator,\n );\n\n let split_options = SplitterChunkHeaderOptions::default();\n let document = RagDocument::new(contents);\n let split_documents = splitter.split_documents(&[document], &split_options);\n rag_files.push(RagFile {\n hash: hash.clone(),\n path,\n documents: split_documents,\n });\n }\n\n let mut next_file_id = self.data.next_file_id;\n let mut files = vec![];\n let mut document_ids = vec![];\n let mut embeddings = vec![];\n\n if !rag_files.is_empty() {\n let mut texts = vec![];\n for file in rag_files.into_iter() {\n for (document_index, document) in file.documents.iter().enumerate() {\n document_ids.push(DocumentId::new(next_file_id, document_index));\n texts.push(document.page_content.clone())\n }\n files.push((next_file_id, file));\n next_file_id += 1;\n }\n\n let embeddings_data = EmbeddingsData::new(texts, false);\n embeddings = self\n .create_embeddings(embeddings_data, spinner.clone())\n .await?;\n }\n\n let to_delete_file_ids: Vec<_> = to_deleted.values().flatten().copied().collect();\n self.data.del(to_delete_file_ids);\n self.data.add(next_file_id, files, document_ids, embeddings);\n self.data.document_paths = document_paths.into_iter().collect();\n\n if self.data.files.is_empty() {\n bail!(\"No RAG files\");\n }\n\n progress(&spinner, \"Building store\".into());\n self.hnsw = self.data.build_hnsw();\n self.bm25 = self.data.build_bm25();\n\n Ok(())\n }\n\n async fn hybird_search(\n &self,\n query: &str,\n top_k: usize,\n rerank_model: Option<&str>,\n ) -> Result> {\n let (vector_search_results, keyword_search_results) = tokio::join!(\n self.vector_search(query, top_k, 0.0),\n self.keyword_search(query, top_k, 0.0),\n );\n\n let vector_search_results = vector_search_results?;\n debug!(\"vector_search_results: {vector_search_results:?}\",);\n let vector_search_ids: Vec =\n vector_search_results.into_iter().map(|(v, _)| v).collect();\n\n let keyword_search_results = keyword_search_results?;\n debug!(\"keyword_search_results: {keyword_search_results:?}\",);\n let keyword_search_ids: Vec =\n keyword_search_results.into_iter().map(|(v, _)| v).collect();\n\n let ids = match rerank_model {\n Some(model_id) => {\n let model =\n Model::retrieve_model(&self.config.read(), model_id, ModelType::Reranker)?;\n let client = init_client(&self.config, Some(model))?;\n let ids: IndexSet = [vector_search_ids, keyword_search_ids]\n .concat()\n .into_iter()\n .collect();\n let mut documents = vec![];\n let mut documents_ids = vec![];\n for id in ids {\n if let Some(document) = self.data.get(id) {\n documents_ids.push(id);\n documents.push(document.page_content.to_string());\n }\n }\n let data = RerankData::new(query.to_string(), documents, top_k);\n let list = client.rerank(&data).await.context(\"Failed to rerank\")?;\n let ids: Vec<_> = list\n .into_iter()\n .take(top_k)\n .filter_map(|item| documents_ids.get(item.index).cloned())\n .collect();\n debug!(\"rerank_ids: {ids:?}\");\n ids\n }\n None => {\n let ids = reciprocal_rank_fusion(\n vec![vector_search_ids, keyword_search_ids],\n vec![1.125, 1.0],\n top_k,\n );\n debug!(\"rrf_ids: {ids:?}\");\n ids\n }\n };\n let output = ids\n .into_iter()\n .filter_map(|id| {\n let document = self.data.get(id)?;\n Some((id, document.page_content.clone()))\n })\n .collect();\n Ok(output)\n }\n\n async fn vector_search(\n &self,\n query: &str,\n top_k: usize,\n min_score: f32,\n ) -> Result> {\n let splitter = RecursiveCharacterTextSplitter::new(\n self.data.chunk_size,\n self.data.chunk_overlap,\n &DEFAULT_SEPARATES,\n );\n let texts = splitter.split_text(query);\n let embeddings_data = EmbeddingsData::new(texts, true);\n let embeddings = self.create_embeddings(embeddings_data, None).await?;\n let output = self\n .hnsw\n .parallel_search(&embeddings, top_k, 30)\n .into_iter()\n .flat_map(|list| {\n list.into_iter()\n .filter_map(|v| {\n let score = 1.0 - v.distance;\n if score > min_score {\n Some((DocumentId(v.d_id), score))\n } else {\n None\n }\n })\n .collect::>()\n })\n .collect();\n Ok(output)\n }\n\n async fn keyword_search(\n &self,\n query: &str,\n top_k: usize,\n min_score: f32,\n ) -> Result> {\n let results = self.bm25.search(query, top_k);\n let output: Vec<(DocumentId, f32)> = results\n .into_iter()\n .filter_map(|v| {\n let score = v.score;\n if score > min_score {\n Some((v.document.id, score))\n } else {\n None\n }\n })\n .collect();\n Ok(output)\n }\n\n async fn create_embeddings(\n &self,\n data: EmbeddingsData,\n spinner: Option,\n ) -> Result {\n let embedding_client = init_client(&self.config, Some(self.embedding_model.clone()))?;\n let EmbeddingsData { texts, query } = data;\n let batch_size = self\n .data\n .batch_size\n .or_else(|| self.embedding_model.max_batch_size());\n let batch_size = match self.embedding_model.max_input_tokens() {\n Some(max_input_tokens) => {\n let x = max_input_tokens / self.data.chunk_size;\n match batch_size {\n Some(y) => x.min(y),\n None => x,\n }\n }\n None => batch_size.unwrap_or(1),\n };\n let mut output = vec![];\n let batch_chunks = texts.chunks(batch_size.max(1));\n let batch_chunks_len = batch_chunks.len();\n let retry_limit = env::var(get_env_name(\"embeddings_retry_limit\"))\n .ok()\n .and_then(|v| v.parse::().ok())\n .unwrap_or(2);\n for (index, texts) in batch_chunks.enumerate() {\n progress(\n &spinner,\n format!(\"Creating embeddings [{}/{batch_chunks_len}]\", index + 1),\n );\n let chunk_data = EmbeddingsData {\n texts: texts.to_vec(),\n query,\n };\n let mut retry = 0;\n let chunk_output = loop {\n retry += 1;\n match embedding_client.embeddings(&chunk_data).await {\n Ok(v) => break v,\n Err(e) if retry < retry_limit => {\n debug!(\"retry {retry} failed: {e}\");\n sleep(Duration::from_secs(2u64.pow(retry - 1))).await;\n continue;\n }\n Err(e) => {\n return Err(e).with_context(|| {\n format!(\"Failed to create embedding after {retry_limit} attempts\")\n })?\n }\n }\n };\n output.extend(chunk_output);\n }\n Ok(output)\n }\n}\n\n#[derive(Clone, Serialize, Deserialize)]\npub struct RagData {\n pub embedding_model: String,\n pub chunk_size: usize,\n pub chunk_overlap: usize,\n pub reranker_model: Option,\n pub top_k: usize,\n pub batch_size: Option,\n pub next_file_id: FileId,\n pub document_paths: Vec,\n pub files: IndexMap,\n #[serde(with = \"serde_vectors\")]\n pub vectors: IndexMap>,\n}\n\nimpl Debug for RagData {\n fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {\n f.debug_struct(\"RagData\")\n .field(\"embedding_model\", &self.embedding_model)\n .field(\"chunk_size\", &self.chunk_size)\n .field(\"chunk_overlap\", &self.chunk_overlap)\n .field(\"reranker_model\", &self.reranker_model)\n .field(\"top_k\", &self.top_k)\n .field(\"batch_size\", &self.batch_size)\n .field(\"next_file_id\", &self.next_file_id)\n .field(\"document_paths\", &self.document_paths)\n .field(\"files\", &self.files)\n .finish()\n }\n}\n\nimpl RagData {\n pub fn new(\n embedding_model: String,\n chunk_size: usize,\n chunk_overlap: usize,\n reranker_model: Option,\n top_k: usize,\n batch_size: Option,\n ) -> Self {\n Self {\n embedding_model,\n chunk_size,\n chunk_overlap,\n reranker_model,\n top_k,\n batch_size,\n next_file_id: 0,\n document_paths: Default::default(),\n files: Default::default(),\n vectors: Default::default(),\n }\n }\n\n pub fn get(&self, id: DocumentId) -> Option<&RagDocument> {\n let (file_index, document_index) = id.split();\n let file = self.files.get(&file_index)?;\n let document = file.documents.get(document_index)?;\n Some(document)\n }\n\n pub fn del(&mut self, file_ids: Vec) {\n for file_id in file_ids {\n if let Some(file) = self.files.swap_remove(&file_id) {\n for (document_index, _) in file.documents.iter().enumerate() {\n let document_id = DocumentId::new(file_id, document_index);\n self.vectors.swap_remove(&document_id);\n }\n }\n }\n }\n\n pub fn add(\n &mut self,\n next_file_id: FileId,\n files: Vec<(FileId, RagFile)>,\n document_ids: Vec,\n embeddings: EmbeddingsOutput,\n ) {\n self.next_file_id = next_file_id;\n self.files.extend(files);\n self.vectors\n .extend(document_ids.into_iter().zip(embeddings));\n }\n\n pub fn build_hnsw(&self) -> Hnsw<'static, f32, DistCosine> {\n let hnsw = Hnsw::new(32, self.vectors.len(), 16, 200, DistCosine {});\n let list: Vec<_> = self.vectors.iter().map(|(k, v)| (v, k.0)).collect();\n hnsw.parallel_insert(&list);\n hnsw\n }\n\n pub fn build_bm25(&self) -> SearchEngine {\n let mut documents = vec![];\n for (file_index, file) in self.files.iter() {\n for (document_index, document) in file.documents.iter().enumerate() {\n let id = DocumentId::new(*file_index, document_index);\n documents.push(bm25::Document::new(id, &document.page_content))\n }\n }\n SearchEngineBuilder::::with_documents(Language::English, documents)\n .k1(1.5)\n .b(0.75)\n .build()\n }\n}\n\n#[derive(Debug, Clone, Serialize, Deserialize)]\npub struct RagFile {\n hash: String,\n path: String,\n documents: Vec,\n}\n\n#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]\npub struct RagDocument {\n pub page_content: String,\n pub metadata: DocumentMetadata,\n}\n\nimpl RagDocument {\n pub fn new>(page_content: S) -> Self {\n RagDocument {\n page_content: page_content.into(),\n metadata: IndexMap::new(),\n }\n }\n}\n\nimpl Default for RagDocument {\n fn default() -> Self {\n RagDocument {\n page_content: \"\".to_string(),\n metadata: IndexMap::new(),\n }\n }\n}\n\npub type FileId = usize;\n\n#[derive(Clone, Copy, Hash, Eq, PartialEq, Ord, PartialOrd)]\npub struct DocumentId(usize);\n\nimpl Debug for DocumentId {\n fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {\n let (file_index, document_index) = self.split();\n f.write_fmt(format_args!(\"{file_index}-{document_index}\"))\n }\n}\n\nimpl DocumentId {\n pub fn new(file_index: usize, document_index: usize) -> Self {\n let value = (file_index << (usize::BITS / 2)) | document_index;\n Self(value)\n }\n\n pub fn split(self) -> (usize, usize) {\n let value = self.0;\n let low_mask = (1 << (usize::BITS / 2)) - 1;\n let low = value & low_mask;\n let high = value >> (usize::BITS / 2);\n (high, low)\n }\n}\n\nfn select_embedding_model(models: &[&Model]) -> Result {\n let models: Vec<_> = models\n .iter()\n .map(|v| SelectOption::new(v.id(), v.description()))\n .collect();\n let result = Select::new(\"Select embedding model:\", models).prompt()?;\n Ok(result.value)\n}\n\n#[derive(Debug)]\nstruct SelectOption {\n pub value: String,\n pub description: String,\n}\n\nimpl SelectOption {\n pub fn new(value: String, description: String) -> Self {\n Self { value, description }\n }\n}\n\nimpl std::fmt::Display for SelectOption {\n fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {\n write!(f, \"{} ({})\", self.value, self.description)\n }\n}\n\nfn set_chunk_size(model: &Model) -> Result {\n let default_value = model.default_chunk_size().to_string();\n let help_message = model\n .max_tokens_per_chunk()\n .map(|v| format!(\"The model's max_tokens is {v}\"));\n\n let mut text = Text::new(\"Set chunk size:\")\n .with_default(&default_value)\n .with_validator(move |text: &str| {\n let out = match text.parse::() {\n Ok(_) => Validation::Valid,\n Err(_) => Validation::Invalid(\"Must be a integer\".into()),\n };\n Ok(out)\n });\n if let Some(help_message) = &help_message {\n text = text.with_help_message(help_message);\n }\n let value = text.prompt()?;\n value.parse().map_err(|_| anyhow!(\"Invalid chunk_size\"))\n}\n\nfn set_chunk_overlay(default_value: usize) -> Result {\n let value = Text::new(\"Set chunk overlay:\")\n .with_default(&default_value.to_string())\n .with_validator(move |text: &str| {\n let out = match text.parse::() {\n Ok(_) => Validation::Valid,\n Err(_) => Validation::Invalid(\"Must be a integer\".into()),\n };\n Ok(out)\n })\n .prompt()?;\n value.parse().map_err(|_| anyhow!(\"Invalid chunk_overlay\"))\n}\n\nfn add_documents() -> Result> {\n let text = Text::new(\"Add documents:\")\n .with_validator(required!(\"This field is required\"))\n .with_help_message(\"e.g. file;dir/;dir/**/*.{md,mdx};loader:resource;url;website/**\")\n .prompt()?;\n let paths = text\n .split(';')\n .filter_map(|v| {\n let v = v.trim().to_string();\n if v.is_empty() {\n None\n } else {\n Some(v)\n }\n })\n .collect();\n Ok(paths)\n}\n\nasync fn resolve_paths>(\n loaders: &HashMap,\n paths: &[T],\n) -> Result<(\n IndexSet,\n IndexSet,\n IndexSet,\n IndexSet,\n IndexSet,\n)> {\n let mut document_paths = IndexSet::new();\n let mut recursive_urls = IndexSet::new();\n let mut urls = IndexSet::new();\n let mut protocol_paths = IndexSet::new();\n let mut absolute_paths = vec![];\n for path in paths {\n let path = path.as_ref().trim();\n if is_url(path) {\n if let Some(start_url) = path.strip_suffix(\"**\") {\n recursive_urls.insert(start_url.to_string());\n } else {\n urls.insert(path.to_string());\n }\n document_paths.insert(path.to_string());\n } else if is_loader_protocol(loaders, path) {\n protocol_paths.insert(path.to_string());\n document_paths.insert(path.to_string());\n } else {\n let resolved_path = resolve_home_dir(path);\n let absolute_path = to_absolute_path(&resolved_path)\n .with_context(|| format!(\"Invalid path '{path}'\"))?;\n absolute_paths.push(resolved_path);\n document_paths.insert(absolute_path);\n }\n }\n let local_paths = expand_glob_paths(&absolute_paths, false).await?;\n Ok((\n document_paths,\n recursive_urls,\n urls,\n protocol_paths,\n local_paths,\n ))\n}\n\nfn progress(spinner: &Option, message: String) {\n if let Some(spinner) = spinner {\n let _ = spinner.set_message(message);\n }\n}\n\nfn reciprocal_rank_fusion(\n list_of_document_ids: Vec>,\n list_of_weights: Vec,\n top_k: usize,\n) -> Vec {\n let rrf_k = top_k * 2;\n let mut map: IndexMap = IndexMap::new();\n for (document_ids, weight) in list_of_document_ids\n .into_iter()\n .zip(list_of_weights.into_iter())\n {\n for (index, &item) in document_ids.iter().enumerate() {\n *map.entry(item).or_default() += (1.0 / ((rrf_k + index + 1) as f32)) * weight;\n }\n }\n let mut sorted_items: Vec<(DocumentId, f32)> = map.into_iter().collect();\n sorted_items.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());\n\n sorted_items\n .into_iter()\n .take(top_k)\n .map(|(v, _)| v)\n .collect()\n}\n" }, { "path": "src/rag/serde_vectors.rs", "content": "use super::*;\n\nuse base64::{engine::general_purpose::STANDARD, Engine};\nuse serde::{de, Deserializer, Serializer};\n\npub fn serialize(\n vectors: &IndexMap>,\n serializer: S,\n) -> Result\nwhere\n S: Serializer,\n{\n let encoded_map: IndexMap = vectors\n .iter()\n .map(|(id, vec)| {\n let (h, l) = id.split();\n let byte_slice = unsafe {\n std::slice::from_raw_parts(\n vec.as_ptr() as *const u8,\n vec.len() * std::mem::size_of::(),\n )\n };\n (format!(\"{h}-{l}\"), STANDARD.encode(byte_slice))\n })\n .collect();\n\n encoded_map.serialize(serializer)\n}\n\npub fn deserialize<'de, D>(deserializer: D) -> Result>, D::Error>\nwhere\n D: Deserializer<'de>,\n{\n let encoded_map: IndexMap =\n IndexMap::::deserialize(deserializer)?;\n\n let mut decoded_map = IndexMap::new();\n for (key, base64_str) in encoded_map {\n let decoded_key: DocumentId = key\n .split_once('-')\n .and_then(|(h, l)| {\n let h = h.parse::().ok()?;\n let l = l.parse::().ok()?;\n Some(DocumentId::new(h, l))\n })\n .ok_or_else(|| de::Error::custom(format!(\"Invalid key '{key}'\")))?;\n\n let decoded_data = STANDARD.decode(&base64_str).map_err(de::Error::custom)?;\n\n if decoded_data.len() % std::mem::size_of::() != 0 {\n return Err(de::Error::custom(format!(\"Invalid vector at '{key}'\")));\n }\n\n let num_f32s = decoded_data.len() / std::mem::size_of::();\n\n let mut vec_f32 = vec![0.0f32; num_f32s];\n unsafe {\n std::ptr::copy_nonoverlapping(\n decoded_data.as_ptr(),\n vec_f32.as_mut_ptr() as *mut u8,\n decoded_data.len(),\n );\n }\n\n decoded_map.insert(decoded_key, vec_f32);\n }\n\n Ok(decoded_map)\n}\n" }, { "path": "src/rag/splitter/language.rs", "content": "#[derive(PartialEq, Eq, Hash)]\npub enum Language {\n Cpp,\n Go,\n Java,\n Js,\n Php,\n Proto,\n Python,\n Rst,\n Ruby,\n Rust,\n Scala,\n Swift,\n Markdown,\n Latex,\n Html,\n Sol,\n}\n\nimpl Language {\n pub fn separators(&self) -> Vec<&str> {\n match self {\n Language::Cpp => vec![\n \"\\nclass \",\n \"\\nvoid \",\n \"\\nint \",\n \"\\nfloat \",\n \"\\ndouble \",\n \"\\nif \",\n \"\\nfor \",\n \"\\nwhile \",\n \"\\nswitch \",\n \"\\ncase \",\n \"\\n\\n\",\n \"\\n\",\n \" \",\n \"\",\n ],\n Language::Go => vec![\n \"\\nfunc \",\n \"\\nvar \",\n \"\\nconst \",\n \"\\ntype \",\n \"\\nif \",\n \"\\nfor \",\n \"\\nswitch \",\n \"\\ncase \",\n \"\\n\\n\",\n \"\\n\",\n \" \",\n \"\",\n ],\n Language::Java => vec![\n \"\\nclass \",\n \"\\npublic \",\n \"\\nprotected \",\n \"\\nprivate \",\n \"\\nstatic \",\n \"\\nif \",\n \"\\nfor \",\n \"\\nwhile \",\n \"\\nswitch \",\n \"\\ncase \",\n \"\\n\\n\",\n \"\\n\",\n \" \",\n \"\",\n ],\n Language::Js => vec![\n \"\\nfunction \",\n \"\\nconst \",\n \"\\nlet \",\n \"\\nvar \",\n \"\\nclass \",\n \"\\nif \",\n \"\\nfor \",\n \"\\nwhile \",\n \"\\nswitch \",\n \"\\ncase \",\n \"\\ndefault \",\n \"\\n\\n\",\n \"\\n\",\n \" \",\n \"\",\n ],\n Language::Php => vec![\n \"\\nfunction \",\n \"\\nclass \",\n \"\\nif \",\n \"\\nforeach \",\n \"\\nwhile \",\n \"\\ndo \",\n \"\\nswitch \",\n \"\\ncase \",\n \"\\n\\n\",\n \"\\n\",\n \" \",\n \"\",\n ],\n Language::Proto => vec![\n \"\\nmessage \",\n \"\\nservice \",\n \"\\nenum \",\n \"\\noption \",\n \"\\nimport \",\n \"\\nsyntax \",\n \"\\n\\n\",\n \"\\n\",\n \" \",\n \"\",\n ],\n Language::Python => vec![\"\\nclass \", \"\\ndef \", \"\\n\\tdef \", \"\\n\\n\", \"\\n\", \" \", \"\"],\n Language::Rst => vec![\n \"\\n===\\n\", \"\\n---\\n\", \"\\n***\\n\", \"\\n.. \", \"\\n\\n\", \"\\n\", \" \", \"\",\n ],\n Language::Ruby => vec![\n \"\\ndef \",\n \"\\nclass \",\n \"\\nif \",\n \"\\nunless \",\n \"\\nwhile \",\n \"\\nfor \",\n \"\\ndo \",\n \"\\nbegin \",\n \"\\nrescue \",\n \"\\n\\n\",\n \"\\n\",\n \" \",\n \"\",\n ],\n Language::Rust => vec![\n \"\\nfn \", \"\\nconst \", \"\\nlet \", \"\\nif \", \"\\nwhile \", \"\\nfor \", \"\\nloop \",\n \"\\nmatch \", \"\\nconst \", \"\\n\\n\", \"\\n\", \" \", \"\",\n ],\n Language::Scala => vec![\n \"\\nclass \",\n \"\\nobject \",\n \"\\ndef \",\n \"\\nval \",\n \"\\nvar \",\n \"\\nif \",\n \"\\nfor \",\n \"\\nwhile \",\n \"\\nmatch \",\n \"\\ncase \",\n \"\\n\\n\",\n \"\\n\",\n \" \",\n \"\",\n ],\n Language::Swift => vec![\n \"\\nfunc \",\n \"\\nclass \",\n \"\\nstruct \",\n \"\\nenum \",\n \"\\nif \",\n \"\\nfor \",\n \"\\nwhile \",\n \"\\ndo \",\n \"\\nswitch \",\n \"\\ncase \",\n \"\\n\\n\",\n \"\\n\",\n \" \",\n \"\",\n ],\n Language::Markdown => vec![\n \"\\n## \",\n \"\\n### \",\n \"\\n#### \",\n \"\\n##### \",\n \"\\n###### \",\n \"```\\n\\n\",\n \"\\n\\n***\\n\\n\",\n \"\\n\\n---\\n\\n\",\n \"\\n\\n___\\n\\n\",\n \"\\n\\n\",\n \"\\n\",\n \" \",\n \"\",\n ],\n Language::Latex => vec![\n \"\\n\\\\chapter{\",\n \"\\n\\\\section{\",\n \"\\n\\\\subsection{\",\n \"\\n\\\\subsubsection{\",\n \"\\n\\\\begin{enumerate}\",\n \"\\n\\\\begin{itemize}\",\n \"\\n\\\\begin{description}\",\n \"\\n\\\\begin{list}\",\n \"\\n\\\\begin{quote}\",\n \"\\n\\\\begin{quotation}\",\n \"\\n\\\\begin{verse}\",\n \"\\n\\\\begin{verbatim}\",\n \"\\n\\\\begin{align}\",\n \"$$\",\n \"$\",\n \"\\n\\n\",\n \"\\n\",\n \" \",\n \"\",\n ],\n Language::Html => vec![\n \"\", \"
\", \"

\", \"
\", \"

  • \", \"

    \", \"

    \", \"

    \", \"

    \", \"

    \",\n \"
    \", \"\", \"\", \"\", \"
    \", \"\", \"
      \", \"
        \", \"
        \",\n \"