apunkt/dograh
1
0
Fork 0
mirror of https://github.com/dograh-hq/dograh.git synced 2026-06-28 08:49:42 +02:00
Code Issues Projects Releases Packages Wiki Activity

fix: enable knowledge base with Dograh config v2

Browse source
This commit is contained in:
Abhishek Kumar 2026-06-25 22:21:11 +05:30
parent d675fd1fda
commit efb25a0cc5
19 changed files with 557 additions and 113 deletions
Download patch file Download diff file Expand all files Collapse all files

43
api/services/workflow/tools/knowledge_base.py
View file

@ -13,8 +13,7 @@ from loguru import logger
from opentelemetry import trace
from api.db import db_client
from api.services.configuration.registry import ServiceProviders
from api.services.gen_ai import AzureOpenAIEmbeddingService, OpenAIEmbeddingService
from api.services.gen_ai import build_embedding_service
from api.services.pipecat.tracing_config import ensure_tracing
@ -266,33 +265,19 @@ async def _perform_retrieval(
"Model Configurations > Embedding."
)
if (
embeddings_provider == ServiceProviders.AZURE.value
and embeddings_endpoint
):
embedding_service = AzureOpenAIEmbeddingService(
db_client=db_client,
api_key=embeddings_api_key,
endpoint=embeddings_endpoint,
model_id=embeddings_model or "text-embedding-3-small",
api_version=embeddings_api_version or "2024-02-15-preview",
)
else:
default_headers = None
if (
embeddings_provider == ServiceProviders.DOGRAH.value
and correlation_id
):
default_headers = {
"X-Dograh-Correlation-Id": correlation_id,
}
embedding_service = OpenAIEmbeddingService(
db_client=db_client,
api_key=embeddings_api_key,
model_id=embeddings_model or "text-embedding-3-small",
base_url=embeddings_base_url,
default_headers=default_headers,
)
# Search runs inside a workflow run: reuse the run's MPS correlation
# id (present only for v2 orgs; None otherwise → sent without the
# protocol). The Dograh-managed path forwards it via request metadata.
embedding_service = await build_embedding_service(
db_client=db_client,
provider=embeddings_provider,
api_key=embeddings_api_key,
model=embeddings_model,
base_url=embeddings_base_url,
endpoint=embeddings_endpoint,
api_version=embeddings_api_version,
correlation_id=correlation_id,
)
results = await embedding_service.search_similar_chunks(
query=query,