class CgRagService:
def __init__(
self,
*,
config: CgRagConfig,
manager: Any,
api_reranker: Any,
args: Any,
capacity_limiters: CapacityLimiters | None = None,
) -> None:
self.config = config
self.manager = manager
self.api_reranker = api_reranker
self.args = args
self.capacity_limiters = capacity_limiters or CapacityLimiters.from_config(config)
self.retrieve_cache = RetrieveCache(CacheConfig(config.cache_max_entries, config.cache_ttl_seconds))
self.generation_session = requests.Session()
self.hydration_service = HydrationService(config=config, args=args, build_on_init=False)
self.hydration_sidecar_error = ""
self.hydration_sidecar = self._build_hydration_sidecar()
if not self.hydration_sidecar_error:
self.hydration_sidecar_error = self.hydration_service.sidecar_error
self.retrieve_service = RetrieveService(self)
self.generation_service = ConstrainedGenerationService(self)
self.stream_service = CgRagStreamService(self)
self.vision_service = CgVisionService(self)
self.health_service = CgRagHealthService(self)
def health(self) -> dict[str, Any]:
return self.health_service.health()
def profiles(self) -> dict[str, Any]:
return self.health_service.profiles()
def retrieve_rerank(
self,
request: CgRetrieveRerankRequest,
*,
cancel_token: Any = None,
) -> dict[str, Any]:
return _call_with_optional_cancel(
self.retrieve_service.retrieve_rerank,
request,
cancel_token=cancel_token,
)
def warmup(self, scopes: list[str] | None = None, *, cancel_token: Any = None) -> dict[str, Any]:
return _call_with_optional_cancel(
self.retrieve_service.warmup,
scopes,
cancel_token=cancel_token,
)
def constrained_generate(
self,
request: CgConstrainedGenerateRequest,
*,
cancel_token: Any = None,
) -> dict[str, Any]:
return _call_with_optional_cancel(
self.generation_service.constrained_generate,
request,
cancel_token=cancel_token,
)
def constrained_generate_stream_events(
self,
request: CgConstrainedGenerateRequest,
*,
cancel_token: Any = None,
) -> Iterable[tuple[str, dict[str, Any]]]:
yield from self.generation_service.constrained_generate_stream_events(
request,
cancel_token=cancel_token,
)
def _answer_token_events(
self,
answer_text: str,
*,
cancel_token: Any = None,
) -> Iterable[tuple[str, dict[str, Any]]]:
yield from self.stream_service.answer_token_events(answer_text, cancel_token=cancel_token)
@staticmethod
def _stream_sleep(delay_seconds: float) -> None:
time.sleep(delay_seconds)
def rag(self, request: CgRagRequest, *, cancel_token: Any = None) -> dict[str, Any]:
return _call_with_optional_cancel(
self.stream_service.rag,
request,
cancel_token=cancel_token,
)
def rag_stream_events(
self,
request: CgRagRequest,
*,
cancel_token: Any = None,
) -> Iterable[tuple[str, dict[str, Any]]]:
yield from self.stream_service.rag_stream_events(request, cancel_token=cancel_token)
def vision_answer(
self,
request: CgVisionAnswerRequest,
*,
cancel_token: Any = None,
) -> dict[str, Any]:
return self.vision_service.vision_answer(request, cancel_token=cancel_token)
def vision_answer_stream_events(
self,
request: CgVisionAnswerRequest,
*,
cancel_token: Any = None,
) -> Iterable[tuple[str, dict[str, Any]]]:
yield from self.vision_service.vision_answer_stream_events(
request, cancel_token=cancel_token
)
def vision_observe(
self,
request: CgVisionAnswerRequest,
*,
cancel_token: Any = None,
) -> dict[str, Any]:
return self.vision_service.vision_observe(request, cancel_token=cancel_token)
def vision_agentic(
self,
request: CgVisionAgenticRequest,
*,
cancel_token: Any = None,
) -> dict[str, Any]:
return self.vision_service.vision_agentic(request, cancel_token=cancel_token)
def vision_agentic_stream_events(
self,
request: CgVisionAgenticRequest,
*,
cancel_token: Any = None,
) -> Iterable[tuple[str, dict[str, Any]]]:
yield from self.vision_service.vision_agentic_stream_events(
request, cancel_token=cancel_token
)
def _api_rerank_enabled(self) -> bool:
return self.retrieve_service.api_rerank_enabled()
def _current_profile_fingerprint(self, manager_health: dict[str, Any]) -> dict[str, Any]:
return self.retrieve_service.current_profile_fingerprint(manager_health)
def _profile_fingerprint(
self,
scope: str,
*,
manager_health: dict[str, Any] | None = None,
) -> dict[str, Any]:
return self.retrieve_service.profile_fingerprint(scope, manager_health=manager_health)
def _profile_for_scope(
self,
scope: str,
*,
manager_health: dict[str, Any] | None = None,
) -> Any:
return self.retrieve_service.profile_for_scope(scope, manager_health=manager_health)
def _retrieval_error(self, request: CgRagRequest) -> dict[str, Any]:
return self.retrieve_service.retrieval_error(request)
@staticmethod
def _parse_chat_completion_sse(lines: Iterable[str | bytes]) -> Iterable[dict[str, Any]]:
yield from ConstrainedGenerationService.parse_chat_completion_sse(lines)
def _generation_thinking_enabled(self, request: CgConstrainedGenerateRequest) -> bool:
return self.generation_service.generation_thinking_enabled(request)
@staticmethod
def _constrained_generation_max_tokens(enable_thinking: bool) -> int:
return ConstrainedGenerationService.constrained_generation_max_tokens(enable_thinking)
def _apply_generation_payload_options(self, payload: dict[str, Any], *, enable_thinking: bool) -> None:
self.generation_service.apply_generation_payload_options(payload, enable_thinking=enable_thinking)
def _post_generation_payload(self, payload: dict[str, Any], *, cancel_token: Any = None) -> dict[str, Any]:
with self.capacity_limiters.generation.acquire(cancel_token=cancel_token):
return self.generation_service.post_generation_payload(payload, cancel_token=cancel_token)
def _post_generation_payload_streaming(
self,
payload: dict[str, Any],
on_token: Callable[[str, str], None] | None = None,
cancel_token: Any = None,
) -> dict[str, Any]:
with self.capacity_limiters.generation.acquire(cancel_token=cancel_token):
return self.generation_service.post_generation_payload_streaming(
payload,
on_token=on_token,
cancel_token=cancel_token,
)
def _generation_headers(self) -> dict[str, str]:
return self.generation_service.generation_headers()
@staticmethod
def _answer_token_chunks(text: str, *, max_chars: int = 36) -> list[str]:
return ConstrainedGenerationService.answer_token_chunks(text, max_chars=max_chars)
def _build_hydration_sidecar(self) -> CorpusHydrationSidecar | None:
return self.hydration_service.build_sidecar()
def _disabled_hydration_health(self) -> dict[str, Any]:
return self.hydration_service.disabled_health(error=self.hydration_sidecar_error)
def _hydration_health(self) -> dict[str, Any]:
return self.hydration_service.health(
sidecar=self.hydration_sidecar,
error=self.hydration_sidecar_error,
)
def _warmup_requested_scopes(self, scopes: list[str] | None) -> list[str]:
return self.retrieve_service.warmup_requested_scopes(scopes)
def _normalize_warmup_scope(self, scope: str | None) -> str:
return self.retrieve_service.normalize_warmup_scope(scope)
def _normalize_manager_scope(self, scope: str | None) -> str:
return self.retrieve_service.normalize_manager_scope(scope)
def _manager_default_scope(self) -> str:
return self.retrieve_service.manager_default_scope()
def _generation_error(self, code: str, message: str) -> dict[str, Any]:
return self.generation_service.generation_error(code, message)
def _parse_generation_indices(self, raw_content: str, max_items: int, max_index: int) -> list[int]:
return self.generation_service.parse_generation_indices(raw_content, max_items, max_index)
def _filter_generation_indices(self, indices: list[Any], max_items: int, max_index: int) -> list[int]:
return self.generation_service.filter_generation_indices(indices, max_items, max_index)
def _extract_response_message(self, body: Any) -> dict[str, Any]:
return self.generation_service.extract_response_message(body)
def _extract_message_content(self, message: dict[str, Any]) -> str:
return self.generation_service.extract_message_content(message)
def _extract_reasoning_content(self, body: dict[str, Any]) -> str:
return self.generation_service.extract_reasoning_content(body)
def _extract_usage(self, body: dict[str, Any]) -> dict[str, int]:
return self.generation_service.extract_usage(body)
def _coerce_token_count(self, value: Any) -> int:
return self.generation_service.coerce_token_count(value)
def _answer_text(self, pred_items: list[str]) -> str:
return self.generation_service.answer_text(pred_items)
def _legacy_profile(self, profile: Any) -> dict[str, Any]:
return self.retrieve_service.legacy_profile(profile)
def _rerank_backend(self, retriever: Any | None) -> str | None:
return self.retrieve_service.rerank_backend(retriever)
def _rerank_model_path(self, retriever: Any | None) -> str:
return self.retrieve_service.rerank_model_path(retriever)
def _embedding_api_status(self, retriever: Any | None) -> dict[str, Any]:
return self.retrieve_service.embedding_api_status(retriever)
def _dense_index_dimension(self, retriever: Any | None) -> int | None:
return self.retrieve_service.dense_index_dimension(retriever)