feat: Add Forge as LLM provider#26
Conversation
## Changes - Automated integration updates from manager loop. Files modified: .env.example README.md functions_for_pipeline.py
📝 WalkthroughWalkthroughReintroduces Changes
Sequence DiagramsequenceDiagram
participant App as Application
participant Helper as Forge-aware helpers
participant ChatLLM as ChatOpenAI
participant Embeddings as OpenAIEmbeddings
participant Forge as Forge API
participant OpenAI as OpenAI API
App->>Helper: create_chat_llm(model_name, **kwargs)
Helper->>Helper: resolve credentials (FORGE_API_KEY, FORGE_API_BASE) or fallback to OPENAI_API_KEY
Helper->>Helper: _forge_model_name(model_name)
alt Forge credentials present
Helper->>ChatLLM: init with Forge base & forged model name
ChatLLM->>Forge: send chat requests
else
Helper->>ChatLLM: init with OpenAI credentials & original model name
ChatLLM->>OpenAI: send chat requests
end
App->>Helper: create_embeddings()
alt Forge credentials present
Helper->>Embeddings: init with "OpenAI/text-embedding-3-small" and Forge creds
Embeddings->>Forge: request embeddings
else
Helper->>Embeddings: init with configured OpenAI embedding model and creds
Embeddings->>OpenAI: request embeddings
end
Estimated Code Review Effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
README.md (1)
116-120: Consider adding brief Forge documentation.The new environment variables are listed, but users may not know what Forge is or when to use it. A brief explanation would help:
Suggested documentation addition
``` OPENAI_API_KEY= GROQ_API_KEY= FORGE_API_KEY= FORGE_API_BASE=https://api.forge.tensorblock.co/v1 ``` + + **Note**: [Forge](https://github.com/tensorblock/forge) is an optional middleware that routes inference across 40+ upstream providers. Set `FORGE_API_KEY` and `FORGE_API_BASE` to use Forge instead of direct OpenAI access. + you can look at the `.env.example` file for reference.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@README.md` around lines 116 - 120, Add a short explanatory note in the README next to the new environment variables (OPENAI_API_KEY, GROQ_API_KEY, FORGE_API_KEY, FORGE_API_BASE) that describes what Forge is and when to use it: say that Forge is an optional middleware that can route inference across 40+ upstream providers, include a link to https://github.com/tensorblock/forge, and instruct users to set FORGE_API_KEY and FORGE_API_BASE when they want to use Forge instead of direct OpenAI access; place this note near the `.env` example so users immediately see the relation to the listed variables.functions_for_pipeline.py (1)
52-55: Hardcoded embedding model may cause inconsistencies.The Forge path hardcodes
"text-embedding-3-small", while the non-Forge path usesOpenAIEmbeddings()which defaults to a different model (typically"text-embedding-ada-002"in older versions). This could cause:
- Dimension mismatch: Different embedding models produce different vector dimensions, which would break existing FAISS indexes.
- No configurability: Users cannot specify a different embedding model when using Forge.
Consider making the model configurable or ensuring consistency:
Proposed fix for consistency
-def create_embeddings(): +def create_embeddings(model: str = "text-embedding-3-small"): if forge_api_key: - return OpenAIEmbeddings(model=_forge_model_name("text-embedding-3-small")) - return OpenAIEmbeddings() + return OpenAIEmbeddings(model=_forge_model_name(model)) + return OpenAIEmbeddings(model=model)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@functions_for_pipeline.py` around lines 52 - 55, The create_embeddings() function currently hardcodes "text-embedding-3-small" for the Forge path and calls OpenAIEmbeddings() with no model for the non-Forge path, risking dimension mismatches; change create_embeddings() to accept or read a single configurable embedding model name (e.g., EMBEDDING_MODEL env/arg) and pass that model into OpenAIEmbeddings for the non-Forge path and to _forge_model_name(EMBEDDING_MODEL) for the Forge path (use the existing forge_api_key, OpenAIEmbeddings, create_embeddings, and _forge_model_name identifiers to locate and update the branches) so both paths use the same configurable model and avoid embedding-dimension inconsistencies.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@functions_for_pipeline.py`:
- Around line 29-36: The current block redundantly re-assigns OPENAI_API_KEY to
itself and also silently overwrites the user's OpenAI credential when
forge_api_key is present; remove the redundant elif that sets
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY") and stop mutating
global OPENAI_* env vars with Forge values. Instead, set/use Forge-specific
environment variables (e.g., FORGE_OPENAI_API_KEY and FORGE_OPENAI_BASE) or pass
forge_api_key/forge_api_base directly into the client constructors where the
symbols forge_api_key and forge_api_base are used, and add a short comment
documenting that Forge creds are separate from the original OPENAI_* variables.
---
Nitpick comments:
In `@functions_for_pipeline.py`:
- Around line 52-55: The create_embeddings() function currently hardcodes
"text-embedding-3-small" for the Forge path and calls OpenAIEmbeddings() with no
model for the non-Forge path, risking dimension mismatches; change
create_embeddings() to accept or read a single configurable embedding model name
(e.g., EMBEDDING_MODEL env/arg) and pass that model into OpenAIEmbeddings for
the non-Forge path and to _forge_model_name(EMBEDDING_MODEL) for the Forge path
(use the existing forge_api_key, OpenAIEmbeddings, create_embeddings, and
_forge_model_name identifiers to locate and update the branches) so both paths
use the same configurable model and avoid embedding-dimension inconsistencies.
In `@README.md`:
- Around line 116-120: Add a short explanatory note in the README next to the
new environment variables (OPENAI_API_KEY, GROQ_API_KEY, FORGE_API_KEY,
FORGE_API_BASE) that describes what Forge is and when to use it: say that Forge
is an optional middleware that can route inference across 40+ upstream
providers, include a link to https://github.com/tensorblock/forge, and instruct
users to set FORGE_API_KEY and FORGE_API_BASE when they want to use Forge
instead of direct OpenAI access; place this note near the `.env` example so
users immediately see the relation to the listed variables.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 32e7c341-fae5-43eb-a24f-870667839736
📒 Files selected for processing (3)
.env.exampleREADME.mdfunctions_for_pipeline.py
When Forge is active, the module was overwriting OPENAI_API_KEY and OPENAI_BASE_URL in the global process environment. This side effect could break other parts of the system that read those env vars. Replace the os.environ mutations with module-level variables (_openai_api_key, _openai_api_base) and pass them explicitly to ChatOpenAI() and OpenAIEmbeddings() via their constructor kwargs. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@functions_for_pipeline.py`:
- Around line 58-66: create_embeddings() switches to
_forge_model_name("text-embedding-3-small") when forge_api_key is present,
causing an embedding-space mismatch with FAISS stores built with the default
OpenAIEmbeddings (text-embedding-ada-002); update create_embeddings (and any
code that loads stores in create_retrievers) so the embedding model is explicit
and consistent: either always set the model to the same string used during
indexing (e.g., "text-embedding-ada-002") or derive the model name from a single
canonical config value used by both indexing and querying (refactor references
to OpenAIEmbeddings, _forge_model_name, and create_retrievers to use that
canonical model); additionally add a runtime check in the store-loading path
(create_retrievers) to detect model vs store metadata mismatch and raise a clear
error instructing to rebuild the store if they differ.
- Around line 64-66: When creating embeddings for Forge in the block that sets
embed_kwargs["model"] = _forge_model_name("text-embedding-3-small") and returns
OpenAIEmbeddings(**embed_kwargs), explicitly set a tiktoken-compatible model
name by adding embed_kwargs["tiktoken_model_name"] (e.g.,
"text-embedding-3-small" or another recognized tiktoken model) or disable
token-length checking by setting
embed_kwargs["check_embedding_ctx_length"]=False before constructing
OpenAIEmbeddings; update the code around embed_kwargs, _forge_model_name, and
the OpenAIEmbeddings(...) call to ensure token counting uses a valid tiktoken
model name or is skipped.
| def create_embeddings(): | ||
| embed_kwargs = {} | ||
| if _openai_api_key: | ||
| embed_kwargs["openai_api_key"] = _openai_api_key | ||
| if _openai_api_base: | ||
| embed_kwargs["openai_api_base"] = _openai_api_base | ||
| if forge_api_key: | ||
| embed_kwargs["model"] = _forge_model_name("text-embedding-3-small") | ||
| return OpenAIEmbeddings(**embed_kwargs) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify how these FAISS stores were built and which embedding model was used.
rg -n -C3 'chunks_vector_store|chapter_summaries_vector_store|book_quotes_vectorstore|FAISS\.(from_|save_local|load_local)|OpenAIEmbeddings\('Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 1838
🏁 Script executed:
#!/bin/bash
# Search for where FAISS stores are created/saved
rg -n 'save_local|from_documents|from_texts' --type py -A 3 -B 3Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 59
🏁 Script executed:
#!/bin/bash
# Look for build, setup, or initialization scripts/files
fd -e py | head -20 | xargs grep -l 'chunks_vector_store\|chapter_summaries_vector_store\|book_quotes_vectorstore' 2>/dev/null || echo "Searching more broadly..."
find . -name '*build*' -o -name '*init*' -o -name '*setup*' 2>/dev/null | head -10Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 102
🏁 Script executed:
#!/bin/bash
# Check the entire functions_for_pipeline.py file for context
wc -l functions_for_pipeline.pyRepository: NirDiamant/Controllable-RAG-Agent
Length of output: 107
🏁 Script executed:
#!/bin/bash
# Search all files (not just Python) for save_local or store creation patterns
rg -n 'save_local|from_documents|from_texts|chunks_vector_store.*=' --type-list | head -5
rg 'save_local|from_documents' -A 2 -B 2Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 466
🏁 Script executed:
#!/bin/bash
# Check for notebook files or other data preparation scripts
find . -name '*.ipynb' -o -name '*build*' -o -name '*prepare*' -o -name '*create*' 2>/dev/nullRepository: NirDiamant/Controllable-RAG-Agent
Length of output: 154
🏁 Script executed:
#!/bin/bash
# Look at the entire create_embeddings and create_retrievers functions with more context
sed -n '58,82p' functions_for_pipeline.pyRepository: NirDiamant/Controllable-RAG-Agent
Length of output: 1273
🏁 Script executed:
#!/bin/bash
# Search for any mentions of forge_api_key to understand its scope
rg -n 'forge_api_key' -B 2 -A 2Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 1298
🏁 Script executed:
#!/bin/bash
# Search in notebook files for where FAISS stores are created
rg 'save_local|from_documents|OpenAIEmbeddings' ./sophisticated_rag_agent_harry_potter.ipynb -A 2 -B 2Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 1965
🏁 Script executed:
#!/bin/bash
# Check what the default embedding model is for OpenAIEmbeddings when no model is specified
# First, let's see what langchain-openai version is used
rg 'openai|langchain' -i requirements.txt setup.py pyproject.toml 2>/dev/null | head -20Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 469
🏁 Script executed:
#!/bin/bash
# Look for any comments or documentation in the code about the embedding models
rg -i 'embedding.*model|text-embedding' -B 2 -A 2Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 408
🏁 Script executed:
#!/bin/bash
# Check if there's any documentation in requirements.txt versions
cat requirements.txtRepository: NirDiamant/Controllable-RAG-Agent
Length of output: 4460
🏁 Script executed:
#!/bin/bash
# Let's check the langchain-openai version to understand the default embedding model
rg 'langchain-openai' requirements.txt -A 1 -B 1Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 133
🏁 Script executed:
#!/bin/bash
# Search for any version pinning or setup.py that might indicate versions
cat setup.py 2>/dev/null || cat pyproject.toml 2>/dev/null || echo "No setup files found"Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 97
🏁 Script executed:
#!/bin/bash
# Check the notebook more carefully to see if there are any clues about when stores were created
rg 'chunks_vector_store|chapter_summaries_vector_store|book_quotes_vectorstore' ./sophisticated_rag_agent_harry_potter.ipynb -B 10 | head -40Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 1968
🏁 Script executed:
#!/bin/bash
# Get more context from functions_for_pipeline.py to understand the full flow
sed -n '1,70p' functions_for_pipeline.py | head -50Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 1656
🏁 Script executed:
#!/bin/bash
# Check if there's any conditional logic around forge_api_key that also affects store creation
rg 'create_embeddings|create_retrievers' -B 5 -A 10Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 2902
Embedding model mismatch when Forge is enabled: FAISS stores were built with the default OpenAI embedding model but are queried with text-embedding-3-small.
The FAISS stores in the repository were created in the notebook (sophisticated_rag_agent_harry_potter.ipynb) using bare OpenAIEmbeddings() calls, which default to text-embedding-ada-002. However, when forge_api_key is set, create_embeddings() at lines 64-65 switches to text-embedding-3-small, causing create_retrievers() to load those stores with incompatible embeddings. This embedding space mismatch will degrade retrieval quality.
Either make the embedding model explicit and consistent on both indexing and querying paths, or add logic to detect model mismatch and require a store rebuild when Forge mode is enabled.
Also applies to: 70-74
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@functions_for_pipeline.py` around lines 58 - 66, create_embeddings() switches
to _forge_model_name("text-embedding-3-small") when forge_api_key is present,
causing an embedding-space mismatch with FAISS stores built with the default
OpenAIEmbeddings (text-embedding-ada-002); update create_embeddings (and any
code that loads stores in create_retrievers) so the embedding model is explicit
and consistent: either always set the model to the same string used during
indexing (e.g., "text-embedding-ada-002") or derive the model name from a single
canonical config value used by both indexing and querying (refactor references
to OpenAIEmbeddings, _forge_model_name, and create_retrievers to use that
canonical model); additionally add a runtime check in the store-loading path
(create_retrievers) to detect model vs store metadata mismatch and raise a clear
error instructing to rebuild the store if they differ.
| if forge_api_key: | ||
| embed_kwargs["model"] = _forge_model_name("text-embedding-3-small") | ||
| return OpenAIEmbeddings(**embed_kwargs) |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
Does langchain_openai.OpenAIEmbeddingsdefaulttiktoken_model_nameto the embeddingmodel, and do the docs recommend overriding tiktoken_model_nameorcheck_embedding_ctx_length for OpenAI-compatible providers with custom model names?
💡 Result:
Yes.
- Defaulting behavior:
langchain_openai.OpenAIEmbeddingsdefinestiktoken_model_name: str | None = None, and its tokenizer path usesmodel_name = self.tiktoken_model_name or self.model, so whentiktoken_model_nameis left asNoneit effectively defaults to the embeddingmodel. [1] - What the docs recommend for OpenAI-compatible providers w/ custom model names:
- The class docstring explicitly recommends, for “OpenAI-compatible APIs (e.g. OpenRouter, Ollama, vLLM)”, setting
check_embedding_ctx_length=Falseso it “send[s] raw text instead of tokens” (i.e., bypasses client-side tokenization/splitting that many compatible providers don’t support). [1] - It also notes you may want to set
tiktoken_model_namewhen your provider’s embeddingmodelname is not supported bytiktoken, to avoid errors whentiktokenis called (e.g., Azure deployments or OpenAI-like providers with different model IDs). [1]
- The class docstring explicitly recommends, for “OpenAI-compatible APIs (e.g. OpenRouter, Ollama, vLLM)”, setting
Sources:
[1] LangChain source: libs/partners/openai/langchain_openai/embeddings/base.py (OpenAIEmbeddings docstring + _tokenize implementation). (raw.githubusercontent.com)
Citations:
🏁 Script executed:
rg "_forge_model_name" -B 2 -A 5Repository: NirDiamant/Controllable-RAG-Agent
Length of output: 1478
Set tiktoken_model_name for Forge embeddings to avoid token counting failures.
_forge_model_name("text-embedding-3-small") produces "OpenAI/text-embedding-3-small". langchain_openai.OpenAIEmbeddings defaults tiktoken_model_name to the embedding model when not explicitly set, and "OpenAI/text-embedding-3-small" is not a recognized tiktoken model, causing token counting to fail before the Forge request is sent. The docs for OpenAIEmbeddings recommend either overriding tiktoken_model_name or disabling check_embedding_ctx_length for OpenAI-compatible providers with custom model names.
🛠️ Proposed fix
def create_embeddings():
embed_kwargs = {}
if _openai_api_key:
embed_kwargs["openai_api_key"] = _openai_api_key
if _openai_api_base:
embed_kwargs["openai_api_base"] = _openai_api_base
if forge_api_key:
embed_kwargs["model"] = _forge_model_name("text-embedding-3-small")
+ embed_kwargs["tiktoken_model_name"] = "text-embedding-3-small"
return OpenAIEmbeddings(**embed_kwargs)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if forge_api_key: | |
| embed_kwargs["model"] = _forge_model_name("text-embedding-3-small") | |
| return OpenAIEmbeddings(**embed_kwargs) | |
| if forge_api_key: | |
| embed_kwargs["model"] = _forge_model_name("text-embedding-3-small") | |
| embed_kwargs["tiktoken_model_name"] = "text-embedding-3-small" | |
| return OpenAIEmbeddings(**embed_kwargs) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@functions_for_pipeline.py` around lines 64 - 66, When creating embeddings for
Forge in the block that sets embed_kwargs["model"] =
_forge_model_name("text-embedding-3-small") and returns
OpenAIEmbeddings(**embed_kwargs), explicitly set a tiktoken-compatible model
name by adding embed_kwargs["tiktoken_model_name"] (e.g.,
"text-embedding-3-small" or another recognized tiktoken model) or disable
token-length checking by setting
embed_kwargs["check_embedding_ctx_length"]=False before constructing
OpenAIEmbeddings; update the code around embed_kwargs, _forge_model_name, and
the OpenAIEmbeddings(...) call to ensure token counting uses a valid tiktoken
model name or is skipped.
1 participant
Summary
Adds Forge API support by introducing
create_chat_llm()andcreate_embeddings()factory functions that route requests through Forge's OpenAI-compatible API whenFORGE_API_KEYis set. Replaces directChatOpenAI()calls with the factory to centralize provider configuration.Changes
.env.example: AddedFORGE_API_KEYandFORGE_API_BASEentriesREADME.md: Documented Forge environment variablesfunctions_for_pipeline.py: Addedcreate_chat_llm(),create_embeddings(), and_forge_model_name()helpers; updated 13 LLM instantiation sites to use the factoryUsage
Test Evidence
All existing tests pass. When
FORGE_API_KEYis not set, behavior is identical to upstream.I work at TensorBlock and will help maintain this integration.
About Forge
Forge is an open-source middleware service for unified AI model provider management. It routes requests across 40+ AI providers with access to thousands of models through a single OpenAI-compatible API.