30 lines
1.2 KiB
Python
30 lines
1.2 KiB
Python
from __future__ import annotations
|
|
|
|
from app.modules.rag.intent_router_v2.models import LayerQuery, RepoContext
|
|
|
|
|
|
class LayerQueryBuilder:
|
|
def build(self, intent: str, repo_context: RepoContext, *, domains: list[str], layers_map: dict[str, list[tuple[str, int]]]) -> list[LayerQuery]:
|
|
available = set(repo_context.available_layers or [])
|
|
result: list[LayerQuery] = []
|
|
for layer_id, top_k in layers_map[intent]:
|
|
if not self._layer_matches_domains(layer_id, domains):
|
|
continue
|
|
if available and layer_id not in available:
|
|
continue
|
|
result.append(LayerQuery(layer_id=layer_id, top_k=top_k))
|
|
if result:
|
|
return result
|
|
return [
|
|
LayerQuery(layer_id=layer_id, top_k=top_k)
|
|
for layer_id, top_k in layers_map[intent]
|
|
if self._layer_matches_domains(layer_id, domains)
|
|
]
|
|
|
|
def _layer_matches_domains(self, layer_id: str, domains: list[str]) -> bool:
|
|
if domains == ["CODE"]:
|
|
return layer_id.startswith("C")
|
|
if domains == ["DOCS"]:
|
|
return layer_id.startswith("D")
|
|
return layer_id.startswith("C") or layer_id.startswith("D")
|