diff --git a/artifacts/enterprise_review.docx b/artifacts/enterprise_review.docx new file mode 100644 index 0000000..502beea Binary files /dev/null and b/artifacts/enterprise_review.docx differ diff --git a/artifacts/enterprise_review.xlsx b/artifacts/enterprise_review.xlsx new file mode 100644 index 0000000..1f20b40 Binary files /dev/null and b/artifacts/enterprise_review.xlsx differ diff --git a/docker-compose.infra.yml b/docker-compose.infra.yml index 11fce9d..8c2c555 100644 --- a/docker-compose.infra.yml +++ b/docker-compose.infra.yml @@ -18,7 +18,7 @@ services: restart: unless-stopped pgadmin: - image: dpage/pgadmin4:8 + image: dpage/pgadmin4:9 container_name: rag-pgadmin environment: PGADMIN_DEFAULT_EMAIL: ${PGADMIN_DEFAULT_EMAIL:-admin@example.com} diff --git a/examples/enterprise_vllm_two_questions.py b/examples/enterprise_vllm_two_questions.py new file mode 100644 index 0000000..8f14b5b --- /dev/null +++ b/examples/enterprise_vllm_two_questions.py @@ -0,0 +1,371 @@ +from __future__ import annotations + +import time +from pathlib import Path + +from ragflow_orchestrator.enterprise import ( + EnterpriseAnswerLLMConfig, + EnterpriseIntentLLMConfig, + EnterprisePipeline, + EnterprisePipelineConfig, + ReviewBundleRequest, +) +from ragflow_orchestrator.enterprise.tools import ToolContext +from ragflow_orchestrator.models import BaseChunk, RetrievalResult + + +class TracingEnterprisePipeline(EnterprisePipeline): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.last_prompt_source: str = "unknown" + self._reset_timing() + + def _reset_timing(self) -> None: + self.timing: dict[str, object] = { + "llm_calls_total": 0, + "llm_time_total_s": 0.0, + "prompt_tokens_total": 0, + "completion_tokens_total": 0, + "tokens_total": 0, + "intent_calls": 0, + "intent_time_s": 0.0, + "answer_calls": 0, + "answer_time_s": 0.0, + "planner_calls": 0, + "planner_time_s": 0.0, + "llm_calls": [], + "review_bundle_time_s": 0.0, + } + + def review_bundle(self, request_payload: ReviewBundleRequest): # type: ignore[override] + self._reset_timing() + started = time.perf_counter() + result = super().review_bundle(request_payload) + self.timing["review_bundle_time_s"] = time.perf_counter() - started + return result + + def _build_prompt(self, req, tasks, grouped, task_outputs, risk_result) -> str: # type: ignore[no-untyped-def] + if self.config.bundle.use_prompt_orchestrator and self._check_prompt_orchestrator_available(): + built = self._build_prompt_via_prompt_orchestrator(req=req, grouped=grouped) + if built: + self.last_prompt_source = "prompt_orchestrator" + return built + + self.last_prompt_source = "fallback" + was_enabled = self.config.bundle.use_prompt_orchestrator + self.config.bundle.use_prompt_orchestrator = False + try: + return super()._build_prompt(req, tasks, grouped, task_outputs, risk_result) + finally: + self.config.bundle.use_prompt_orchestrator = was_enabled + + def _call_llm_text(self, *, prompt: str, cfg, custom_client): # type: ignore[override,no-untyped-def] + started = time.perf_counter() + result = super()._call_llm_text(prompt=prompt, cfg=cfg, custom_client=custom_client) + duration = time.perf_counter() - started + + provider = str(getattr(cfg, "provider", "unknown")) + model = str(getattr(cfg, "model", "unknown")) + meta = getattr(self, "_last_llm_meta", {}) + prompt_tokens = int(meta.get("prompt_tokens", 0)) if isinstance(meta, dict) else 0 + completion_tokens = int(meta.get("completion_tokens", 0)) if isinstance(meta, dict) else 0 + total_tokens = int(meta.get("total_tokens", 0)) if isinstance(meta, dict) else 0 + endpoint = str(meta.get("endpoint", "unknown")) if isinstance(meta, dict) else "unknown" + prompt_preview = prompt[:120].replace("\n", " ") + phase = "answer" + lowered = prompt.lower() + if "analyze enterprise review intent" in lowered: + phase = "intent" + self.timing["intent_calls"] = int(self.timing["intent_calls"]) + 1 + self.timing["intent_time_s"] = float(self.timing["intent_time_s"]) + duration + elif "planning tool calls for answer generation" in lowered: + phase = "planner" + self.timing["planner_calls"] = int(self.timing["planner_calls"]) + 1 + self.timing["planner_time_s"] = float(self.timing["planner_time_s"]) + duration + else: + self.timing["answer_calls"] = int(self.timing["answer_calls"]) + 1 + self.timing["answer_time_s"] = float(self.timing["answer_time_s"]) + duration + + self.timing["llm_calls_total"] = int(self.timing["llm_calls_total"]) + 1 + self.timing["llm_time_total_s"] = float(self.timing["llm_time_total_s"]) + duration + self.timing["prompt_tokens_total"] = int(self.timing["prompt_tokens_total"]) + prompt_tokens + self.timing["completion_tokens_total"] = int(self.timing["completion_tokens_total"]) + completion_tokens + self.timing["tokens_total"] = int(self.timing["tokens_total"]) + total_tokens + llm_calls = self.timing["llm_calls"] + assert isinstance(llm_calls, list) + llm_calls.append( + { + "phase": phase, + "provider": provider, + "model": model, + "endpoint": endpoint, + "duration_s": round(duration, 3), + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": total_tokens, + "prompt_preview": prompt_preview, + } + ) + return result + + def _generate_answer(self, prompt: str, task_outputs, req): # type: ignore[override,no-untyped-def] + started = time.perf_counter() + result = super()._generate_answer(prompt, task_outputs, req) + duration = time.perf_counter() - started + + cfg = self.config.answer_llm + if str(getattr(cfg, "provider", "none")) == "none": + return result + + provider = str(getattr(cfg, "provider", "unknown")) + model = str(getattr(cfg, "model", "unknown")) + meta = getattr(self, "_last_llm_meta", {}) + prompt_tokens = int(meta.get("prompt_tokens", 0)) if isinstance(meta, dict) else 0 + completion_tokens = int(meta.get("completion_tokens", 0)) if isinstance(meta, dict) else 0 + total_tokens = int(meta.get("total_tokens", 0)) if isinstance(meta, dict) else 0 + endpoint = str(meta.get("endpoint", "unknown")) if isinstance(meta, dict) else "unknown" + + self.timing["answer_calls"] = int(self.timing["answer_calls"]) + 1 + self.timing["answer_time_s"] = float(self.timing["answer_time_s"]) + duration + self.timing["llm_calls_total"] = int(self.timing["llm_calls_total"]) + 1 + self.timing["llm_time_total_s"] = float(self.timing["llm_time_total_s"]) + duration + self.timing["prompt_tokens_total"] = int(self.timing["prompt_tokens_total"]) + prompt_tokens + self.timing["completion_tokens_total"] = int(self.timing["completion_tokens_total"]) + completion_tokens + self.timing["tokens_total"] = int(self.timing["tokens_total"]) + total_tokens + + llm_calls = self.timing["llm_calls"] + assert isinstance(llm_calls, list) + llm_calls.append( + { + "phase": "answer", + "provider": provider, + "model": model, + "endpoint": endpoint, + "duration_s": round(duration, 3), + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": total_tokens, + "prompt_preview": prompt[:120].replace("\n", " "), + } + ) + return result + + +class StaticDemoRetriever: + def __init__(self) -> None: + self._rows = [ + RetrievalResult( + chunk=BaseChunk( + id="c-001", + text=( + "All electrical equipment must include grounding and emergency stop controls. " + "Periodic inspection interval is every 6 months." + ), + metadata={ + "source_type": "normative", + "document_id": "doc-safety-01", + "version_id": "v3", + "page": "12", + "clause_path": "5.2.3", + "valid_to": "2099-12-31", + }, + source_id="doc-safety-01", + chunk_index=0, + ), + score=0.94, + provider="static-demo", + ), + RetrievalResult( + chunk=BaseChunk( + id="c-002", + text=( + "Operator training is mandatory before commissioning. " + "A signed checklist must be stored in the project archive." + ), + metadata={ + "source_type": "process", + "document_id": "doc-safety-02", + "version_id": "v1", + "page": "4", + "clause_path": "2.1", + "valid_to": "2099-12-31", + }, + source_id="doc-safety-02", + chunk_index=0, + ), + score=0.89, + provider="static-demo", + ), + ] + + def retrieve( + self, + question: str, + top_k: int = 5, + filters: dict[str, object] | None = None, + ) -> list[RetrievalResult]: + _ = question, filters + return self._rows[:top_k] + + +def build_pipeline() -> TracingEnterprisePipeline: + config = EnterprisePipelineConfig( + intent_llm=EnterpriseIntentLLMConfig( + provider="openai", + model="Qwen/Qwen3-32B-AWQ", + openai_base_url="http://hppii:8000/v1", + openai_api_key="EMPTY", + temperature=0.0, + max_tokens=180, + enable_tools=False, + allowed_tools=["retrieve_context", "group_evidence", "score_risks", "build_attribution"], + max_tool_calls=3, + ), + answer_llm=EnterpriseAnswerLLMConfig( + provider="openai", + model="Qwen/Qwen3-32B-AWQ", + openai_base_url="http://hppii:8000/v1", + openai_api_key="EMPTY", + temperature=0.1, + max_tokens=900, + enable_tools=False, + allowed_tools=[ + "retrieve_context", + "group_evidence", + "score_risks", + "build_attribution", + "create_word_document", + "create_xlsx_table", + ], + max_tool_calls=5, + ), + ) + return TracingEnterprisePipeline(query_retriever=StaticDemoRetriever(), config=config) + + +def main() -> None: + pipeline = build_pipeline() + + questions = [ + "Выдели ключевые требования безопасности и дай краткое резюме.", + "Найди потенциальные риски несоответствия и предложи приоритетные действия.", + ] + + results = [] + prompt_sources: list[str] = [] + timings: list[dict[str, object]] = [] + for idx, question in enumerate(questions, start=1): + req = ReviewBundleRequest( + session_id=f"vllm-demo-{idx}", + query_text=question, + document_ids=["doc-safety-01", "doc-safety-02"], + department_principals=["safety-team", "qa-team"], + ) + result = pipeline.review_bundle(req) + results.append(result) + prompt_sources.append(pipeline.last_prompt_source) + timings.append(dict(pipeline.timing)) + + output_dir = Path("artifacts") + output_dir.mkdir(parents=True, exist_ok=True) + + report_lines = [] + for idx, item in enumerate(results, start=1): + report_lines.append(f"Question {idx}: {questions[idx - 1]}") + report_lines.append("Answer:") + report_lines.append(item.llm_answer) + report_lines.append("") + + word_res = pipeline.answer_tools.call( + name="create_word_document", + context=ToolContext(request_id="vllm-demo-artifacts"), + args={ + "filename": str(output_dir / "enterprise_review.docx"), + "title": "Enterprise Review (vLLM Demo)", + "paragraphs": report_lines, + }, + ) + + xlsx_rows: list[list[str]] = [] + for idx, item in enumerate(results, start=1): + xlsx_rows.append( + [ + str(idx), + questions[idx - 1], + " | ".join(task.value for task in item.executed_tasks), + item.llm_answer[:500], + ] + ) + + xlsx_res = pipeline.answer_tools.call( + name="create_xlsx_table", + context=ToolContext(request_id="vllm-demo-artifacts"), + args={ + "filename": str(output_dir / "enterprise_review.xlsx"), + "sheet_name": "Answers", + "headers": ["QuestionNo", "Question", "ExecutedTasks", "AnswerPreview"], + "rows": xlsx_rows, + }, + ) + + print("Processed questions:") + for idx, q in enumerate(questions, start=1): + print(f"{idx}. {q}") + + print("\nPrompt diagnostics:") + for idx, item in enumerate(results, start=1): + source = prompt_sources[idx - 1] + timing = timings[idx - 1] + print(f"Question {idx}: uses PromptOrchestrator = {source == 'prompt_orchestrator'}") + print(f"Prompt source: {source}") + print( + "Timing: " + f"review_bundle={float(timing['review_bundle_time_s']):.3f}s, " + f"llm_total={float(timing['llm_time_total_s']):.3f}s, " + f"llm_calls={int(timing['llm_calls_total'])}, " + f"intent={float(timing['intent_time_s']):.3f}s/{int(timing['intent_calls'])} calls, " + f"planner={float(timing['planner_time_s']):.3f}s/{int(timing['planner_calls'])} calls, " + f"answer={float(timing['answer_time_s']):.3f}s/{int(timing['answer_calls'])} calls, " + f"tokens={int(timing['tokens_total'])} " + f"(prompt={int(timing['prompt_tokens_total'])}, completion={int(timing['completion_tokens_total'])})" + ) + print("LLM calls:") + llm_calls = timing["llm_calls"] + assert isinstance(llm_calls, list) + for call in llm_calls: + if not isinstance(call, dict): + continue + print( + f" - phase={call.get('phase')} provider={call.get('provider')} " + f"model={call.get('model')} endpoint={call.get('endpoint')} " + f"duration={call.get('duration_s')}s tokens={call.get('total_tokens')} " + f"(prompt={call.get('prompt_tokens')}, completion={call.get('completion_tokens')})" + ) + print("Generated prompt:") + print(f"===== PROMPT_BEGIN_Q{idx} =====") + print(item.prompt) + print(f"===== PROMPT_END_Q{idx} =====") + print("-" * 100) + + total_review = sum(float(t["review_bundle_time_s"]) for t in timings) + total_llm = sum(float(t["llm_time_total_s"]) for t in timings) + total_calls = sum(int(t["llm_calls_total"]) for t in timings) + total_prompt_tokens = sum(int(t["prompt_tokens_total"]) for t in timings) + total_completion_tokens = sum(int(t["completion_tokens_total"]) for t in timings) + total_tokens = sum(int(t["tokens_total"]) for t in timings) + print("\nTiming summary:") + print(f"Total review_bundle time: {total_review:.3f}s") + print(f"Total LLM time: {total_llm:.3f}s") + print(f"Total LLM calls: {total_calls}") + print( + f"Total tokens: {total_tokens} " + f"(prompt={total_prompt_tokens}, completion={total_completion_tokens})" + ) + + print("\nArtifacts:") + print(f"DOCX -> {word_res}") + print(f"XLSX -> {xlsx_res}") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 2338979..e6e7fbc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "ragflow_orchestrator" -version = "0.1.16" +version = "0.1.18" description = "Universal RAG orchestration layer with pluggable vector DB adapters" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" @@ -27,6 +27,7 @@ classifiers = [ ] dependencies = [ "mypy>=2.1.0", + "openai>=1.0.0", "psycopg[binary]>=3.1", "pydantic>=2.8,<3", "ruff>=0.15.19", diff --git a/requirements.txt b/requirements.txt index 1b250e1..aadcae0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,7 @@ pytest>=8.2 ruff>=0.5.0 mypy>=1.10 pylint>=3.2 +openai>=1.0.0 psycopg[binary]>=3.1 beautifulsoup4>=4.12 jusText>=3.0.1 diff --git a/src/ragflow_orchestrator/__init__.py b/src/ragflow_orchestrator/__init__.py index 999d89e..ca8e5a3 100644 --- a/src/ragflow_orchestrator/__init__.py +++ b/src/ragflow_orchestrator/__init__.py @@ -1,6 +1,17 @@ from ragflow_orchestrator.config import ConfigStore, EmbeddingConfig, ModuleConfig, PipelineConfig, ProviderConfig from ragflow_orchestrator.context import DocChunk from ragflow_orchestrator.embedding import HashEmbedder, HFEmbedder, OllamaEmbedder, create_embedder +from ragflow_orchestrator.enterprise import ( + ACLAwareRetrieverAdapter, + EnterpriseAnswerLLMConfig, + EnterpriseIntentLLMConfig, + EnterprisePipeline, + EnterprisePipelineConfig, + ReviewBundleRequest, + ReviewBundleResult, + ReviewTask, + StaticRetrieverAdapter, +) from ragflow_orchestrator.evaluation import ( DatasetItem, RerankProfileConfig, @@ -96,6 +107,15 @@ "StrategyEvaluationReport", "evaluate_strategies", "evaluate_rerank_profiles", + "EnterprisePipeline", + "ACLAwareRetrieverAdapter", + "StaticRetrieverAdapter", + "EnterpriseIntentLLMConfig", + "EnterpriseAnswerLLMConfig", + "EnterprisePipelineConfig", + "ReviewBundleRequest", + "ReviewBundleResult", + "ReviewTask", "LanguageMode", "TemplatesConfig", "TemplateRunMetrics", diff --git a/src/ragflow_orchestrator/enterprise/__init__.py b/src/ragflow_orchestrator/enterprise/__init__.py new file mode 100644 index 0000000..338642b --- /dev/null +++ b/src/ragflow_orchestrator/enterprise/__init__.py @@ -0,0 +1,44 @@ +from ragflow_orchestrator.enterprise.adapters import ACLAwareRetrieverAdapter, StaticRetrieverAdapter +from ragflow_orchestrator.enterprise.config import ( + EnterpriseAnswerLLMConfig, + EnterpriseBundleConfig, + EnterpriseIntentLLMConfig, + EnterpriseLLMConfig, + EnterprisePipelineConfig, + EnterpriseRiskConfig, + EnterpriseToolingConfig, +) +from ragflow_orchestrator.enterprise.models import ( + EvidenceItem, + ReviewBundleRequest, + ReviewBundleResult, + ReviewTask, + RiskAssessmentResult, + RiskEntry, + TaskOutput, +) +from ragflow_orchestrator.enterprise.pipeline import EnterprisePipeline +from ragflow_orchestrator.enterprise.tools import EnterpriseToolRegistry, MCPToolProvider, ToolContext + +__all__ = [ + "ACLAwareRetrieverAdapter", + "StaticRetrieverAdapter", + "EnterpriseAnswerLLMConfig", + "EnterpriseBundleConfig", + "EnterpriseIntentLLMConfig", + "EnterpriseLLMConfig", + "EnterprisePipelineConfig", + "EnterpriseRiskConfig", + "EnterpriseToolingConfig", + "EvidenceItem", + "ReviewBundleRequest", + "ReviewBundleResult", + "ReviewTask", + "RiskAssessmentResult", + "RiskEntry", + "TaskOutput", + "EnterprisePipeline", + "EnterpriseToolRegistry", + "MCPToolProvider", + "ToolContext", +] diff --git a/src/ragflow_orchestrator/enterprise/adapters.py b/src/ragflow_orchestrator/enterprise/adapters.py new file mode 100644 index 0000000..711945d --- /dev/null +++ b/src/ragflow_orchestrator/enterprise/adapters.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from ragflow_orchestrator.models import RetrievalQuery, RetrievalResult +from ragflow_orchestrator.protocols import Embedder, RAGProvider + + +class ACLAwareRetrieverAdapter: + """Query retriever adapter with explicit ACL propagation into RetrievalQuery.""" + + def __init__( + self, + *, + provider: RAGProvider, + embedder: Embedder, + include_deleted: bool = False, + ) -> None: + self.provider = provider + self.embedder = embedder + self.include_deleted = include_deleted + + def retrieve( + self, + question: str, + top_k: int = 5, + filters: dict[str, object] | None = None, + ) -> list[RetrievalResult]: + payload = dict(filters or {}) + acl_principals_raw = payload.pop("acl_principals", None) + include_deleted_raw = payload.pop("include_deleted", self.include_deleted) + + acl_principals: list[str] | None = None + if isinstance(acl_principals_raw, list): + acl_principals = [str(item) for item in acl_principals_raw] + + include_deleted = bool(include_deleted_raw) + + query = RetrievalQuery( + text_query=question, + top_k=top_k, + dense_vector=self.embedder.embed(question), + filters=[{"key": key, "value": value} for key, value in payload.items()], + include_deleted=include_deleted, + acl_principals=acl_principals, + ) + return self.provider.retrieve(query) + + +class StaticRetrieverAdapter: + """Small adapter useful for tests and offline simulations.""" + + def __init__(self, rows: list[RetrievalResult]) -> None: + self._rows = rows + + def retrieve( + self, + question: str, + top_k: int = 5, + filters: dict[str, object] | None = None, + ) -> list[RetrievalResult]: + _ = question, filters + return self._rows[:top_k] diff --git a/src/ragflow_orchestrator/enterprise/config.py b/src/ragflow_orchestrator/enterprise/config.py new file mode 100644 index 0000000..2384cd5 --- /dev/null +++ b/src/ragflow_orchestrator/enterprise/config.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, Field + + +class EnterpriseLLMConfig(BaseModel): + provider: Literal["none", "openai", "ollama", "custom"] = "none" + model: str = "gpt-4o-mini" + max_tokens: int = 1000 + temperature: float = 0.1 + enable_tools: bool = False + max_tool_calls: int = 3 + allowed_tools: list[str] = Field(default_factory=list) + openai_api_key: str | None = None + openai_base_url: str | None = None + ollama_base_url: str = "http://localhost:11434" + timeout_seconds: int = 60 + + +class EnterpriseIntentLLMConfig(EnterpriseLLMConfig): + model: str = "gpt-4o-mini" + + +class EnterpriseAnswerLLMConfig(EnterpriseLLMConfig): + model: str = "gpt-4o-mini" + + +class EnterpriseRiskConfig(BaseModel): + enabled_by_default: bool = False + severity_weight: float = 0.5 + likelihood_weight: float = 0.3 + impact_weight: float = 0.2 + + +class EnterpriseToolingConfig(BaseModel): + enable_tools: bool = True + enable_mcp: bool = True + strict_tool_failures: bool = False + + +class EnterpriseBundleConfig(BaseModel): + default_tasks: list[str] = Field( + default_factory=lambda: [ + "requirements_extraction", + "compliance_check", + "consistency_validation", + "completeness_gap_analysis", + "ambiguity_precision_review", + ] + ) + max_evidence_total: int = 24 + max_evidence_per_group: int = 8 + top_k_per_retrieval: int = 8 + group_by_metadata_key: str = "source_type" + use_prompt_orchestrator: bool = True + + +class EnterprisePipelineConfig(BaseModel): + intent_llm: EnterpriseIntentLLMConfig = Field(default_factory=EnterpriseIntentLLMConfig) + answer_llm: EnterpriseAnswerLLMConfig = Field(default_factory=EnterpriseAnswerLLMConfig) + risk: EnterpriseRiskConfig = Field(default_factory=EnterpriseRiskConfig) + tooling: EnterpriseToolingConfig = Field(default_factory=EnterpriseToolingConfig) + bundle: EnterpriseBundleConfig = Field(default_factory=EnterpriseBundleConfig) diff --git a/src/ragflow_orchestrator/enterprise/models.py b/src/ragflow_orchestrator/enterprise/models.py new file mode 100644 index 0000000..1f32616 --- /dev/null +++ b/src/ragflow_orchestrator/enterprise/models.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from datetime import date +from enum import Enum + +from pydantic import BaseModel, Field + + +class ReviewTask(str, Enum): + REQUIREMENTS_EXTRACTION = "requirements_extraction" + COMPLIANCE_CHECK = "compliance_check" + CONSISTENCY_VALIDATION = "consistency_validation" + COMPLETENESS_GAP_ANALYSIS = "completeness_gap_analysis" + AMBIGUITY_PRECISION_REVIEW = "ambiguity_precision_review" + PROJECT_ERROR_ANALYSIS = "project_error_analysis" + CHANGE_IMPACT_DIFF = "change_impact_diff" + RISK_ASSESSMENT = "risk_assessment" + + +class ReviewBundleRequest(BaseModel): + session_id: str + query_text: str + document_ids: list[str] = Field(default_factory=list) + department_principals: list[str] = Field(default_factory=list) + as_of_date: date | None = None + requested_tasks: list[ReviewTask] = Field(default_factory=list) + enable_risk_assessment: bool | None = None + extra_context: dict[str, str] = Field(default_factory=dict) + + +class EvidenceItem(BaseModel): + id: str + group: str + content: str + score: float | None = None + attribution: str = "" + metadata: dict[str, str] = Field(default_factory=dict) + + +class TaskOutput(BaseModel): + task: ReviewTask + summary: str + findings: list[str] = Field(default_factory=list) + evidence_ids: list[str] = Field(default_factory=list) + + +class RiskEntry(BaseModel): + title: str + severity: str + likelihood: float + impact: float + confidence: float + priority_score: float + mitigation: str + evidence_ids: list[str] = Field(default_factory=list) + + +class RiskAssessmentResult(BaseModel): + enabled: bool + entries: list[RiskEntry] = Field(default_factory=list) + + +class ReviewBundleResult(BaseModel): + requested_tasks: list[ReviewTask] + executed_tasks: list[ReviewTask] + grouped_evidence: dict[str, list[EvidenceItem]] + task_outputs: list[TaskOutput] + risk_assessment: RiskAssessmentResult + prompt: str + llm_answer: str diff --git a/src/ragflow_orchestrator/enterprise/pipeline.py b/src/ragflow_orchestrator/enterprise/pipeline.py new file mode 100644 index 0000000..5d7b672 --- /dev/null +++ b/src/ragflow_orchestrator/enterprise/pipeline.py @@ -0,0 +1,923 @@ +from __future__ import annotations + +import importlib +import importlib.util +import json +from collections import defaultdict +from collections.abc import Sequence +from datetime import date, datetime +from pathlib import Path +from typing import Any, Protocol, cast +from urllib import request as urllib_request + +from ragflow_orchestrator.models import RetrievalResult +from ragflow_orchestrator.orchestrator import RAGOrchestrator + +from .adapters import ACLAwareRetrieverAdapter +from .config import EnterpriseLLMConfig, EnterprisePipelineConfig +from .models import ( + EvidenceItem, + ReviewBundleRequest, + ReviewBundleResult, + ReviewTask, + RiskAssessmentResult, + RiskEntry, + TaskOutput, +) +from .tools import EnterpriseToolRegistry, MCPToolProvider, ToolContext + + +class QueryRetriever(Protocol): + def retrieve( + self, + question: str, + top_k: int = 5, + filters: dict[str, object] | None = None, + ) -> list[RetrievalResult]: + ... + + +class EnterprisePipeline: + """Adaptive enterprise review bundle for document-centric RAG workflows.""" + + def __init__( + self, + *, + query_retriever: QueryRetriever, + config: EnterprisePipelineConfig | None = None, + mcp_provider: MCPToolProvider | None = None, + llm_client: object | None = None, + intent_llm_client: object | None = None, + answer_llm_client: object | None = None, + ) -> None: + self.query_retriever = query_retriever + self.config = config or EnterprisePipelineConfig() + self._mcp_provider = mcp_provider + self._intent_llm_client = intent_llm_client or llm_client + self._answer_llm_client = answer_llm_client or llm_client + self._prompt_orchestrator_available: bool | None = None + self._last_llm_meta: dict[str, Any] = {} + + self.tools = EnterpriseToolRegistry( + mcp_provider=mcp_provider, + enable_mcp=self.config.tooling.enable_mcp, + strict_failures=self.config.tooling.strict_tool_failures, + ) + self.intent_tools = EnterpriseToolRegistry( + mcp_provider=mcp_provider, + enable_mcp=self.config.tooling.enable_mcp, + strict_failures=self.config.tooling.strict_tool_failures, + ) + self.answer_tools = EnterpriseToolRegistry( + mcp_provider=mcp_provider, + enable_mcp=self.config.tooling.enable_mcp, + strict_failures=self.config.tooling.strict_tool_failures, + ) + self._register_default_tools() + + @classmethod + def from_orchestrator( + cls, + orchestrator: RAGOrchestrator, + *, + config: EnterprisePipelineConfig | None = None, + mcp_provider: MCPToolProvider | None = None, + llm_client: object | None = None, + intent_llm_client: object | None = None, + answer_llm_client: object | None = None, + ) -> EnterprisePipeline: + retriever = ACLAwareRetrieverAdapter( + provider=orchestrator.provider, + embedder=orchestrator.embedder, + ) + return cls( + query_retriever=retriever, + config=config, + mcp_provider=mcp_provider, + llm_client=llm_client, + intent_llm_client=intent_llm_client, + answer_llm_client=answer_llm_client, + ) + + def set_config(self, config: EnterprisePipelineConfig) -> None: + self.config = config + self.tools = EnterpriseToolRegistry( + mcp_provider=self._mcp_provider, + enable_mcp=self.config.tooling.enable_mcp, + strict_failures=self.config.tooling.strict_tool_failures, + ) + self.intent_tools = EnterpriseToolRegistry( + mcp_provider=self._mcp_provider, + enable_mcp=self.config.tooling.enable_mcp, + strict_failures=self.config.tooling.strict_tool_failures, + ) + self.answer_tools = EnterpriseToolRegistry( + mcp_provider=self._mcp_provider, + enable_mcp=self.config.tooling.enable_mcp, + strict_failures=self.config.tooling.strict_tool_failures, + ) + self._register_default_tools() + + def review_bundle(self, request_payload: ReviewBundleRequest) -> ReviewBundleResult: + tasks = self._resolve_tasks(request_payload) + evidence = self._retrieve_evidence(request_payload) + grouped = self._group_evidence(evidence) + + task_outputs = [self._execute_task(task, request_payload, grouped) for task in tasks] + + risk_enabled = ( + request_payload.enable_risk_assessment + if request_payload.enable_risk_assessment is not None + else self.config.risk.enabled_by_default + ) + risk_result = self._run_risk_assessment(task_outputs, grouped, enabled=risk_enabled) + + prompt = self._build_prompt(request_payload, tasks, grouped, task_outputs, risk_result) + llm_answer = self._generate_answer(prompt, task_outputs, request_payload) + + return ReviewBundleResult( + requested_tasks=tasks, + executed_tasks=tasks, + grouped_evidence=grouped, + task_outputs=task_outputs, + risk_assessment=risk_result, + prompt=prompt, + llm_answer=llm_answer, + ) + + def call_tool(self, name: str, request_id: str, args: dict[str, Any]) -> dict[str, Any]: + return self.tools.call(name=name, context=ToolContext(request_id=request_id), args=args) + + # --------------------------------------------------------------------- + # Internal orchestration + # --------------------------------------------------------------------- + def _resolve_tasks(self, req: ReviewBundleRequest) -> list[ReviewTask]: + if req.requested_tasks: + tasks = list(req.requested_tasks) + else: + llm_tasks = self._analyze_query_with_llm(req) + if llm_tasks: + tasks = llm_tasks + else: + tasks = [ReviewTask(item) for item in self.config.bundle.default_tasks] + lowered = req.query_text.lower() + if any(token in lowered for token in ("риск", "risk", "опас", "угроз")): + tasks.append(ReviewTask.RISK_ASSESSMENT) + if any(token in lowered for token in ("ошиб", "дефект", "problem", "issue")): + tasks.append(ReviewTask.PROJECT_ERROR_ANALYSIS) + if any(token in lowered for token in ("противореч", "consisten", "conflict")): + tasks.append(ReviewTask.CONSISTENCY_VALIDATION) + + seen: set[ReviewTask] = set() + unique: list[ReviewTask] = [] + for task in tasks: + if task in seen: + continue + seen.add(task) + unique.append(task) + return unique + + def _retrieve_evidence(self, req: ReviewBundleRequest) -> list[EvidenceItem]: + base_filters = self._compose_filters(req) + rows = self.query_retriever.retrieve( + question=req.query_text, + top_k=self.config.bundle.top_k_per_retrieval, + filters=base_filters, + ) + + out: list[EvidenceItem] = [] + for row in rows: + chunk = row.chunk + metadata_raw = chunk.metadata if isinstance(chunk.metadata, dict) else {} + metadata = {str(k): str(v) for k, v in metadata_raw.items()} + + if not self._is_document_allowed(req.document_ids, metadata): + continue + if not self._is_valid_on_date(req.as_of_date, metadata): + continue + + group = metadata.get(self.config.bundle.group_by_metadata_key) or "ungrouped" + out.append( + EvidenceItem( + id=chunk.id, + group=group, + content=chunk.text, + score=row.score, + attribution=self._build_attribution(metadata, chunk.id), + metadata=metadata, + ) + ) + + if len(out) > self.config.bundle.max_evidence_total: + out = sorted(out, key=lambda item: item.score or 0.0, reverse=True)[: self.config.bundle.max_evidence_total] + return out + + def _compose_filters(self, req: ReviewBundleRequest) -> dict[str, object]: + filters: dict[str, object] = {} + if req.department_principals: + filters["acl_principals"] = req.department_principals + filters.update(req.extra_context) + return filters + + def _group_evidence(self, evidence: list[EvidenceItem]) -> dict[str, list[EvidenceItem]]: + grouped: dict[str, list[EvidenceItem]] = defaultdict(list) + for item in evidence: + grouped[item.group].append(item) + for key in grouped: + grouped[key] = sorted(grouped[key], key=lambda x: x.score or 0.0, reverse=True)[: self.config.bundle.max_evidence_per_group] + return dict(grouped) + + def _execute_task( + self, + task: ReviewTask, + req: ReviewBundleRequest, + grouped: dict[str, list[EvidenceItem]], + ) -> TaskOutput: + findings: list[str] = [] + evidence_ids: list[str] = [] + + total = sum(len(items) for items in grouped.values()) + if total == 0: + return TaskOutput( + task=task, + summary="No evidence available for this task.", + findings=["No context chunks found after ACL/validity filtering."], + evidence_ids=[], + ) + + for group_name, items in grouped.items(): + if not items: + continue + evidence_ids.extend(item.id for item in items) + top = items[0] + if task == ReviewTask.REQUIREMENTS_EXTRACTION: + findings.append(f"[{group_name}] Candidate requirement sources identified ({len(items)} chunks).") + elif task == ReviewTask.COMPLIANCE_CHECK: + findings.append(f"[{group_name}] Compliance signals detected; top evidence score={top.score or 0.0:.3f}.") + elif task == ReviewTask.CONSISTENCY_VALIDATION: + findings.append(f"[{group_name}] Consistency check candidates extracted.") + elif task == ReviewTask.COMPLETENESS_GAP_ANALYSIS: + findings.append(f"[{group_name}] Potential coverage gaps to review manually.") + elif task == ReviewTask.AMBIGUITY_PRECISION_REVIEW: + findings.append(f"[{group_name}] Ambiguous or weakly specified statements may exist.") + elif task == ReviewTask.PROJECT_ERROR_ANALYSIS: + findings.append(f"[{group_name}] Potential project mistakes flagged for review.") + elif task == ReviewTask.CHANGE_IMPACT_DIFF: + findings.append(f"[{group_name}] Impact analysis requires baseline/version comparison input.") + elif task == ReviewTask.RISK_ASSESSMENT: + findings.append(f"[{group_name}] Risk indicators collected for scoring.") + + summary = f"Task {task.value} processed with {len(evidence_ids)} evidence items." + return TaskOutput(task=task, summary=summary, findings=findings, evidence_ids=sorted(set(evidence_ids))) + + def _run_risk_assessment( + self, + task_outputs: list[TaskOutput], + grouped: dict[str, list[EvidenceItem]], + *, + enabled: bool, + ) -> RiskAssessmentResult: + if not enabled: + return RiskAssessmentResult(enabled=False, entries=[]) + + risk_signals = 0 + evidence_ids: list[str] = [] + for task in task_outputs: + if task.task in { + ReviewTask.COMPLIANCE_CHECK, + ReviewTask.CONSISTENCY_VALIDATION, + ReviewTask.PROJECT_ERROR_ANALYSIS, + ReviewTask.COMPLETENESS_GAP_ANALYSIS, + }: + risk_signals += max(1, len(task.findings)) + evidence_ids.extend(task.evidence_ids) + + severity = "low" + if risk_signals >= 8: + severity = "high" + elif risk_signals >= 4: + severity = "medium" + + likelihood = min(1.0, 0.2 + (risk_signals * 0.08)) + impact = min(1.0, 0.3 + (len(grouped) * 0.1)) + confidence = min(1.0, 0.5 + (len(evidence_ids) * 0.01)) + + score = ( + self.config.risk.severity_weight * (1.0 if severity == "high" else (0.7 if severity == "medium" else 0.4)) + + self.config.risk.likelihood_weight * likelihood + + self.config.risk.impact_weight * impact + ) + + entry = RiskEntry( + title="Composite document review risk", + severity=severity, + likelihood=round(likelihood, 3), + impact=round(impact, 3), + confidence=round(confidence, 3), + priority_score=round(score, 3), + mitigation="Review top mismatches and update missing/ambiguous requirements before approval.", + evidence_ids=sorted(set(evidence_ids))[:20], + ) + return RiskAssessmentResult(enabled=True, entries=[entry]) + + def _build_prompt( + self, + req: ReviewBundleRequest, + tasks: list[ReviewTask], + grouped: dict[str, list[EvidenceItem]], + task_outputs: list[TaskOutput], + risk_result: RiskAssessmentResult, + ) -> str: + if self.config.bundle.use_prompt_orchestrator and self._check_prompt_orchestrator_available(): + built = self._build_prompt_via_prompt_orchestrator(req=req, grouped=grouped) + if built: + return built + + lines: list[str] = [] + lines.append("You are an enterprise document reviewer. Use evidence only.") + lines.append(f"User request: {req.query_text}") + lines.append(f"Tasks: {', '.join(task.value for task in tasks)}") + lines.append("") + + lines.append("Evidence by group:") + for group_name, items in grouped.items(): + lines.append(f"- Group: {group_name} ({len(items)} items)") + for item in items: + score = f"{item.score:.3f}" if item.score is not None else "n/a" + lines.append(f" - [{item.id}] score={score} :: {item.attribution}") + lines.append(f" {item.content[:400]}") + + lines.append("") + lines.append("Intermediate findings:") + for out in task_outputs: + lines.append(f"- {out.task.value}: {out.summary}") + for finding in out.findings[:4]: + lines.append(f" - {finding}") + + if risk_result.enabled and risk_result.entries: + lines.append("") + lines.append("Risk assessment:") + for risk in risk_result.entries: + lines.append( + f"- {risk.title}: severity={risk.severity}, likelihood={risk.likelihood}, impact={risk.impact}, priority={risk.priority_score}" + ) + + lines.append("") + lines.append("Return a structured answer with sections: Summary, Findings, Risks, Recommended actions, Citations.") + return "\n".join(lines) + + def _build_prompt_via_prompt_orchestrator( + self, + *, + req: ReviewBundleRequest, + grouped: dict[str, list[EvidenceItem]], + ) -> str | None: + try: + prompt_module = importlib.import_module("prompt_orchestrator") + state_module = importlib.import_module("prompt_orchestrator.context.state") + except Exception: + return None + + PromptBuilder = getattr(prompt_module, "PromptBuilder", None) + PromptConfig = getattr(prompt_module, "PromptConfig", None) + PromptContextState = getattr(prompt_module, "PromptContextState", None) + DocChunk = getattr(state_module, "DocChunk", None) + if not all((PromptBuilder, PromptConfig, PromptContextState, DocChunk)): + return None + + PromptBuilderCls = cast(type[Any], PromptBuilder) + PromptConfigCls = cast(type[Any], PromptConfig) + PromptContextStateCls = cast(type[Any], PromptContextState) + DocChunkCls = cast(type[Any], DocChunk) + + rag_chunks: list[Any] = [] + for group_name, items in grouped.items(): + for item in items: + rag_chunks.append( + DocChunkCls( + id=item.id, + content=f"[{group_name}] {item.content}\nSource: {item.attribution}", + score=item.score, + metadata=item.metadata, + ) + ) + + cfg = PromptConfigCls( + system_prompt="You are an enterprise document reviewer.", + role="Compliance and Risk Analyst", + task="Use only provided evidence and include citations for every substantial claim.", + constraints=[ + "Do not invent facts.", + "Always include citation IDs in findings.", + "Highlight uncertainty if evidence is insufficient.", + ], + output_format="Markdown", + examples=[], + ) + + state = PromptContextStateCls(session_id=req.session_id, summary=None, recent_messages=[], rag_chunks=rag_chunks) + builder = PromptBuilderCls() + sections = builder.build_sections(config=cfg, state=state, user_message=req.query_text, include_headers=False) + return "\n\n".join([sections["static"], sections["summary"], sections["recent"], sections["rag"]]) + + def _generate_answer(self, prompt: str, task_outputs: list[TaskOutput], req: ReviewBundleRequest) -> str: + cfg = self.config.answer_llm + provider = cfg.provider + if provider == "none": + merged = [f"{item.task.value}: {item.summary}" for item in task_outputs] + return "\n".join(merged) if merged else "No findings." + + tool_results: list[dict[str, Any]] = [] + if cfg.enable_tools: + tool_results = self._plan_and_execute_answer_tools(req=req, base_prompt=prompt) + if tool_results: + prompt = f"{prompt}\n\nTool execution results:\n{json.dumps(tool_results, ensure_ascii=False, indent=2)}" + + if provider == "custom": + client = self._answer_llm_client + if client is None: + return "Custom LLM provider selected, but no client was supplied." + text = client.generate( + prompt=prompt, + model=cfg.model, + max_tokens=cfg.max_tokens, + temperature=cfg.temperature, + ) + return str(text).strip() + + if provider == "openai": + return self._generate_openai(prompt, cfg) + if provider == "ollama": + return self._generate_ollama(prompt, cfg) + + return "Unsupported LLM provider configuration." + + # --------------------------------------------------------------------- + # LLM providers + # --------------------------------------------------------------------- + def _generate_openai(self, prompt: str, cfg: EnterpriseLLMConfig) -> str: + self._last_llm_meta = {} + try: + from openai import OpenAI # type: ignore[import-not-found] + except ImportError: + return "OpenAI client is not installed." + + http_client = None + try: + import httpx # type: ignore[import-not-found] + + # Avoid inherited proxy settings that can break local/private vLLM endpoints. + http_client = httpx.Client(timeout=cfg.timeout_seconds, trust_env=False) + except Exception: + http_client = None + + client = OpenAI( + api_key=cfg.openai_api_key, + base_url=cfg.openai_base_url, + http_client=http_client, + ) + + try: + chat = client.chat.completions.create( + model=cfg.model, + messages=[{"role": "user", "content": prompt}], + temperature=cfg.temperature, + max_tokens=cfg.max_tokens, + ) + chat_usage = getattr(chat, "usage", None) + self._last_llm_meta = { + "endpoint": "chat.completions", + "prompt_tokens": int(getattr(chat_usage, "prompt_tokens", 0) or 0), + "completion_tokens": int(getattr(chat_usage, "completion_tokens", 0) or 0), + "total_tokens": int(getattr(chat_usage, "total_tokens", 0) or 0), + } + return str(chat.choices[0].message.content or "").strip() + except Exception: + pass + + response = client.responses.create( + model=cfg.model, + input=prompt, + temperature=cfg.temperature, + max_output_tokens=cfg.max_tokens, + ) + usage_obj = getattr(response, "usage", None) + self._last_llm_meta = { + "endpoint": "responses", + "prompt_tokens": int(getattr(usage_obj, "prompt_tokens", 0) or 0), + "completion_tokens": int(getattr(usage_obj, "completion_tokens", 0) or 0), + "total_tokens": int(getattr(usage_obj, "total_tokens", 0) or 0), + } + text = getattr(response, "output_text", None) + if text: + return str(text).strip() + return "" + + def _generate_ollama(self, prompt: str, cfg: EnterpriseLLMConfig) -> str: + endpoint = f"{cfg.ollama_base_url.rstrip('/')}/api/generate" + payload = { + "model": cfg.model, + "prompt": prompt, + "stream": False, + "options": { + "temperature": cfg.temperature, + "num_predict": cfg.max_tokens, + }, + } + req = urllib_request.Request( + endpoint, + data=json.dumps(payload).encode("utf-8"), + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urllib_request.urlopen(req, timeout=cfg.timeout_seconds) as resp: + body = resp.read().decode("utf-8") + parsed = json.loads(body) + return str(parsed.get("response", "")).strip() + except Exception as exc: + return f"Ollama generation failed: {exc}" + + def _analyze_query_with_llm(self, req: ReviewBundleRequest) -> list[ReviewTask] | None: + cfg = self.config.intent_llm + if cfg.provider == "none": + return None + + available = [task.value for task in ReviewTask] + prompt = ( + "Analyze enterprise review intent and return strict JSON with keys: " + "tasks (list[str]), tool_calls (list[object]), notes (string). " + f"Allowed tasks: {available}. " + "Return JSON only.\n" + f"User query: {req.query_text}\n" + f"Extra context: {json.dumps(req.extra_context, ensure_ascii=False)}" + ) + + raw = self._call_llm_text(prompt=prompt, cfg=cfg, custom_client=self._intent_llm_client) + parsed = self._safe_json_load(raw) + tool_calls = self._extract_tool_calls(parsed) + if cfg.enable_tools and tool_calls: + tool_results = self._execute_tool_calls( + registry=self.intent_tools, + request_id=req.session_id, + tool_calls=tool_calls, + cfg=cfg, + ) + retry_prompt = ( + f"{prompt}\n\nTool results:\n{json.dumps(tool_results, ensure_ascii=False, indent=2)}\n" + "Using tool results, return final strict JSON only with keys tasks, tool_calls, notes." + ) + raw = self._call_llm_text(prompt=retry_prompt, cfg=cfg, custom_client=self._intent_llm_client) + parsed = self._safe_json_load(raw) + + if not isinstance(parsed, dict): + return None + + candidate = parsed.get("tasks") + if not isinstance(candidate, list): + return None + + resolved: list[ReviewTask] = [] + for item in candidate: + try: + resolved.append(ReviewTask(str(item))) + except Exception: + continue + return resolved or None + + def _plan_and_execute_answer_tools(self, req: ReviewBundleRequest, base_prompt: str) -> list[dict[str, Any]]: + cfg = self.config.answer_llm + available = self.answer_tools.names() + if not available: + return [] + + planner_prompt = ( + "You are planning tool calls for answer generation. Return strict JSON with key tool_calls only. " + "Each item: {name: str, args: object}. Return JSON only.\n" + f"Available tools: {available}\n" + f"User request: {req.query_text}\n" + f"Prompt summary: {base_prompt[:1200]}" + ) + raw = self._call_llm_text(prompt=planner_prompt, cfg=cfg, custom_client=self._answer_llm_client) + parsed = self._safe_json_load(raw) + tool_calls = self._extract_tool_calls(parsed) + if not tool_calls: + return [] + return self._execute_tool_calls( + registry=self.answer_tools, + request_id=req.session_id, + tool_calls=tool_calls, + cfg=cfg, + ) + + def _call_llm_text(self, *, prompt: str, cfg: EnterpriseLLMConfig, custom_client: object | None) -> str: + if cfg.provider == "custom": + if custom_client is None: + return "" + text = custom_client.generate( + prompt=prompt, + model=cfg.model, + max_tokens=cfg.max_tokens, + temperature=cfg.temperature, + ) + return str(text).strip() + if cfg.provider == "openai": + return self._generate_openai(prompt, cfg) + if cfg.provider == "ollama": + return self._generate_ollama(prompt, cfg) + return "" + + @staticmethod + def _safe_json_load(raw: str) -> dict[str, Any] | list[Any] | None: + text = (raw or "").strip() + if not text: + return None + try: + return json.loads(text) + except Exception: + pass + + start = text.find("{") + end = text.rfind("}") + if start >= 0 and end > start: + try: + return json.loads(text[start : end + 1]) + except Exception: + return None + return None + + @staticmethod + def _extract_tool_calls(payload: dict[str, Any] | list[Any] | None) -> list[dict[str, Any]]: + if not isinstance(payload, dict): + return [] + calls = payload.get("tool_calls") + if not isinstance(calls, list): + return [] + + out: list[dict[str, Any]] = [] + for call in calls: + if not isinstance(call, dict): + continue + name = call.get("name") + args = call.get("args") + if not isinstance(name, str): + continue + if not isinstance(args, dict): + args = {} + out.append({"name": name, "args": args}) + return out + + def _execute_tool_calls( + self, + *, + registry: EnterpriseToolRegistry, + request_id: str, + tool_calls: Sequence[dict[str, Any]], + cfg: EnterpriseLLMConfig, + ) -> list[dict[str, Any]]: + allowed = set(cfg.allowed_tools) + limited = list(tool_calls)[: max(0, cfg.max_tool_calls)] + out: list[dict[str, Any]] = [] + for call in limited: + name = str(call.get("name") or "") + args = call.get("args") + if not isinstance(args, dict): + args = {} + if allowed and name not in allowed: + out.append({"ok": False, "tool": name, "error": "Tool is not in allowed_tools"}) + continue + res = registry.call(name=name, context=ToolContext(request_id=request_id), args=args) + out.append({"tool": name, "result": res}) + return out + + # --------------------------------------------------------------------- + # Utility filters and formatting + # --------------------------------------------------------------------- + @staticmethod + def _build_attribution(metadata: dict[str, str], chunk_id: str) -> str: + doc = metadata.get("document_id") or metadata.get("doc_id") or "unknown-doc" + version = metadata.get("version_id") or metadata.get("version") or "unknown-version" + page = metadata.get("page") or metadata.get("page_number") or "?" + clause = metadata.get("clause_path") or metadata.get("standard_ref") or "" + valid_to = metadata.get("valid_to") or "" + parts = [f"doc={doc}", f"version={version}", f"page={page}", f"chunk={chunk_id}"] + if clause: + parts.append(f"clause={clause}") + if valid_to: + parts.append(f"valid_to={valid_to}") + return ", ".join(parts) + + @staticmethod + def _is_document_allowed(document_ids: list[str], metadata: dict[str, str]) -> bool: + if not document_ids: + return True + doc_id = metadata.get("document_id") or metadata.get("doc_id") + source_id = metadata.get("source_id") + return bool(doc_id and doc_id in document_ids) or bool(source_id and source_id in document_ids) + + @staticmethod + def _is_valid_on_date(as_of_date: date | None, metadata: dict[str, str]) -> bool: + if as_of_date is None: + return True + raw = (metadata.get("valid_to") or "").strip() + if not raw: + return True + + # Accept common date strings and ISO datetime. + candidates = [raw] + if "T" in raw: + candidates.append(raw.split("T", 1)[0]) + + for item in candidates: + try: + valid_to = datetime.fromisoformat(item).date() + return valid_to >= as_of_date + except Exception: + continue + return True + + def _check_prompt_orchestrator_available(self) -> bool: + if self._prompt_orchestrator_available is not None: + return self._prompt_orchestrator_available + try: + self._prompt_orchestrator_available = importlib.util.find_spec("prompt_orchestrator") is not None + except Exception: + self._prompt_orchestrator_available = False + return self._prompt_orchestrator_available + + def _register_default_tools(self) -> None: + if not self.config.tooling.enable_tools: + return + + self.tools.register("retrieve_context", self._tool_retrieve_context) + self.tools.register("group_evidence", self._tool_group_evidence) + self.tools.register("score_risks", self._tool_score_risks) + self.tools.register("build_attribution", self._tool_build_attribution) + + self.intent_tools.register("retrieve_context", self._tool_retrieve_context) + self.intent_tools.register("group_evidence", self._tool_group_evidence) + self.intent_tools.register("score_risks", self._tool_score_risks) + self.intent_tools.register("build_attribution", self._tool_build_attribution) + + self.answer_tools.register("retrieve_context", self._tool_retrieve_context) + self.answer_tools.register("group_evidence", self._tool_group_evidence) + self.answer_tools.register("score_risks", self._tool_score_risks) + self.answer_tools.register("build_attribution", self._tool_build_attribution) + self.answer_tools.register("create_word_document", self._tool_create_word_document) + self.answer_tools.register("create_xlsx_table", self._tool_create_xlsx_table) + + # --------------------------------------------------------------------- + # Default tools + # --------------------------------------------------------------------- + def _tool_retrieve_context(self, context: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + question = str(args.get("question") or "") + top_k = int(args.get("top_k") or self.config.bundle.top_k_per_retrieval) + filters = args.get("filters") + if not isinstance(filters, dict): + filters = None + + rows = self.query_retriever.retrieve(question=question, top_k=top_k, filters=filters) + return { + "ok": True, + "request_id": context.request_id, + "count": len(rows), + "items": [ + { + "id": row.chunk.id, + "score": row.score, + "content": row.chunk.text, + "metadata": row.chunk.metadata, + } + for row in rows + ], + } + + def _tool_group_evidence(self, context: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + items = args.get("items") + key = str(args.get("key") or self.config.bundle.group_by_metadata_key) + if not isinstance(items, list): + return {"ok": False, "request_id": context.request_id, "error": "items must be a list"} + + grouped: dict[str, int] = defaultdict(int) + for item in items: + if not isinstance(item, dict): + continue + metadata = item.get("metadata") + if isinstance(metadata, dict): + group = str(metadata.get(key) or "ungrouped") + else: + group = "ungrouped" + grouped[group] += 1 + + return {"ok": True, "request_id": context.request_id, "groups": dict(grouped)} + + def _tool_score_risks(self, context: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + severity = str(args.get("severity") or "medium").lower() + likelihood = float(args.get("likelihood") or 0.5) + impact = float(args.get("impact") or 0.5) + severity_value = 1.0 if severity == "high" else (0.7 if severity == "medium" else 0.4) + score = ( + self.config.risk.severity_weight * severity_value + + self.config.risk.likelihood_weight * likelihood + + self.config.risk.impact_weight * impact + ) + return { + "ok": True, + "request_id": context.request_id, + "priority_score": round(score, 3), + } + + def _tool_build_attribution(self, context: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + metadata = args.get("metadata") + chunk_id = str(args.get("chunk_id") or "unknown") + if not isinstance(metadata, dict): + return {"ok": False, "request_id": context.request_id, "error": "metadata must be a dict"} + + text = self._build_attribution({str(k): str(v) for k, v in metadata.items()}, chunk_id) + return {"ok": True, "request_id": context.request_id, "attribution": text} + + def _tool_create_word_document(self, context: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + filename = str(args.get("filename") or f"review-{context.request_id}.docx") + title = str(args.get("title") or "Review Report") + content = str(args.get("content") or "") + paragraphs_raw = args.get("paragraphs") + paragraphs: list[str] = [] + if isinstance(paragraphs_raw, list): + paragraphs = [str(item) for item in paragraphs_raw] + + try: + from docx import Document # type: ignore[import-not-found] + except ImportError: + return { + "ok": False, + "request_id": context.request_id, + "error": "python-docx is not installed", + } + + path = Path(filename) + if not path.is_absolute(): + path = Path.cwd() / path + path.parent.mkdir(parents=True, exist_ok=True) + + doc = Document() + doc.add_heading(title, level=1) + if content: + doc.add_paragraph(content) + for paragraph in paragraphs: + doc.add_paragraph(paragraph) + doc.save(str(path)) + + return { + "ok": True, + "request_id": context.request_id, + "path": str(path), + "bytes": path.stat().st_size, + } + + def _tool_create_xlsx_table(self, context: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + filename = str(args.get("filename") or f"review-{context.request_id}.xlsx") + sheet_name = str(args.get("sheet_name") or "Review") + headers_raw = args.get("headers") + rows_raw = args.get("rows") + + headers: list[str] = [] + if isinstance(headers_raw, list): + headers = [str(item) for item in headers_raw] + + rows: list[list[str]] = [] + if isinstance(rows_raw, list): + for row in rows_raw: + if isinstance(row, list): + rows.append([str(item) for item in row]) + + try: + from openpyxl import Workbook # type: ignore[import-not-found] + except ImportError: + return { + "ok": False, + "request_id": context.request_id, + "error": "openpyxl is not installed", + } + + path = Path(filename) + if not path.is_absolute(): + path = Path.cwd() / path + path.parent.mkdir(parents=True, exist_ok=True) + + workbook = Workbook() + worksheet = workbook.active + worksheet.title = sheet_name[:31] or "Review" + + if headers: + worksheet.append(headers) + for row in rows: + worksheet.append(row) + + workbook.save(str(path)) + return { + "ok": True, + "request_id": context.request_id, + "path": str(path), + "rows": len(rows), + "bytes": path.stat().st_size, + } diff --git a/src/ragflow_orchestrator/enterprise/tools.py b/src/ragflow_orchestrator/enterprise/tools.py new file mode 100644 index 0000000..9e63394 --- /dev/null +++ b/src/ragflow_orchestrator/enterprise/tools.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Protocol + + +class MCPToolProvider(Protocol): + def list_tools(self) -> list[str]: + ... + + def call_tool(self, name: str, args: dict[str, Any]) -> dict[str, Any]: + ... + + +@dataclass(slots=True) +class ToolContext: + request_id: str + + +ToolCallable = Callable[[ToolContext, dict[str, Any]], dict[str, Any]] + + +class EnterpriseToolRegistry: + def __init__( + self, + *, + mcp_provider: MCPToolProvider | None = None, + enable_mcp: bool = True, + strict_failures: bool = False, + ) -> None: + self._tools: dict[str, ToolCallable] = {} + self._mcp_provider = mcp_provider + self._enable_mcp = enable_mcp + self._strict_failures = strict_failures + + def register(self, name: str, tool: ToolCallable) -> None: + self._tools[name] = tool + + def names(self) -> list[str]: + names = list(self._tools) + if self._mcp_provider is not None and self._enable_mcp: + try: + names.extend(self._mcp_provider.list_tools()) + except Exception: + if self._strict_failures: + raise + return sorted(set(names)) + + def call(self, name: str, context: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + local = self._tools.get(name) + if local is not None: + try: + return local(context, args) + except Exception as exc: + if self._strict_failures: + raise + return {"ok": False, "error": str(exc), "tool": name} + + if self._mcp_provider is not None and self._enable_mcp: + try: + return self._mcp_provider.call_tool(name, args) + except Exception as exc: + if self._strict_failures: + raise + return {"ok": False, "error": str(exc), "tool": name} + + return {"ok": False, "error": f"Tool '{name}' is not registered", "tool": name} diff --git a/tests/test_enterprise_pipeline.py b/tests/test_enterprise_pipeline.py new file mode 100644 index 0000000..d90de09 --- /dev/null +++ b/tests/test_enterprise_pipeline.py @@ -0,0 +1,212 @@ +from __future__ import annotations + +from datetime import date + +from ragflow_orchestrator.enterprise import ( + ACLAwareRetrieverAdapter, + EnterpriseAnswerLLMConfig, + EnterpriseIntentLLMConfig, + EnterprisePipeline, + EnterprisePipelineConfig, + ReviewBundleRequest, + ReviewTask, +) +from ragflow_orchestrator.models import BaseChunk, RetrievalResult +from ragflow_orchestrator.orchestrator import RAGOrchestrator +from ragflow_orchestrator.presets import document_preset + + +class DummyRetriever: + def retrieve(self, question: str, top_k: int = 5, filters: dict[str, object] | None = None) -> list[RetrievalResult]: + _ = question, top_k, filters + chunk = BaseChunk( + id="c1", + text="Document clause says equipment must have grounding and emergency stop.", + metadata={ + "source_type": "normative", + "document_id": "doc-1", + "version_id": "v1", + "page": "12", + "clause_path": "5.2.3", + "valid_to": "2099-12-31", + }, + source_id="doc-1", + chunk_index=0, + ) + return [RetrievalResult(chunk=chunk, score=0.91, provider="dummy")] + + +class DummyProvider: + name = "dummy" + + def ensure_schema(self, vector_dim: int) -> None: + _ = vector_dim + + def upsert_chunks(self, chunks: list[BaseChunk]) -> None: + _ = chunks + + def delete_chunks(self, chunk_ids: list[str], soft_delete: bool = True) -> None: + _ = chunk_ids, soft_delete + + def retrieve(self, query): + _ = query + chunk = BaseChunk(id="d1", text="dummy", metadata={}, source_id="d", chunk_index=0) + return [RetrievalResult(chunk=chunk, score=0.5, provider=self.name)] + + def healthcheck(self) -> bool: + return True + + +class DummyEmbedder: + dimensions = 4 + + def embed(self, text: str) -> list[float]: + _ = text + return [0.1, 0.2, 0.3, 0.4] + + def embed_many(self, texts): + return [self.embed(text) for text in texts] + + +class IntentLLMStub: + def generate(self, prompt: str, model: str, max_tokens: int, temperature: float) -> str: + _ = prompt, model, max_tokens, temperature + return '{"tasks": ["risk_assessment", "project_error_analysis"], "tool_calls": [], "notes": "ok"}' + + +class AnswerLLMStub: + def __init__(self) -> None: + self.prompts: list[str] = [] + + def generate(self, prompt: str, model: str, max_tokens: int, temperature: float) -> str: + _ = model, max_tokens, temperature + self.prompts.append(prompt) + if "planning tool calls for answer generation" in prompt.lower(): + return ( + '{"tool_calls": [' + '{"name": "score_risks", "args": {"severity": "high", "likelihood": 0.9, "impact": 0.8}}' + "]}" + ) + return "Generated answer" + + +def test_review_bundle_runs_with_defaults() -> None: + pipeline = EnterprisePipeline(query_retriever=DummyRetriever()) + result = pipeline.review_bundle( + ReviewBundleRequest( + session_id="s1", + query_text="Check compliance and find possible issues", + department_principals=["dept_a"], + as_of_date=date.today(), + document_ids=["doc-1"], + ) + ) + + assert result.executed_tasks + assert result.grouped_evidence + assert result.prompt + assert result.llm_answer + + +def test_review_bundle_respects_optional_risk_flag() -> None: + pipeline = EnterprisePipeline(query_retriever=DummyRetriever()) + result = pipeline.review_bundle( + ReviewBundleRequest( + session_id="s2", + query_text="Assess requirements", + requested_tasks=[ReviewTask.REQUIREMENTS_EXTRACTION], + enable_risk_assessment=False, + ) + ) + + assert result.risk_assessment.enabled is False + assert result.risk_assessment.entries == [] + + +def test_acl_aware_adapter_forwards_acl_principals() -> None: + captured = {} + + class CapturingProvider(DummyProvider): + def retrieve(self, query): + captured["acl_principals"] = query.acl_principals + return super().retrieve(query) + + adapter = ACLAwareRetrieverAdapter(provider=CapturingProvider(), embedder=DummyEmbedder()) + adapter.retrieve("hello", top_k=3, filters={"acl_principals": ["dept_a", "dept_b"], "source_type": "normative"}) + + assert captured["acl_principals"] == ["dept_a", "dept_b"] + + +def test_from_orchestrator_constructor_and_tool_call() -> None: + preset = document_preset() + orchestrator = RAGOrchestrator( + provider=DummyProvider(), + embedder=DummyEmbedder(), + chunker=preset.chunker, + cleaner=preset.cleaner, + ) + pipeline = EnterprisePipeline.from_orchestrator( + orchestrator, + config=EnterprisePipelineConfig(), + ) + + response = pipeline.call_tool( + name="score_risks", + request_id="req-1", + args={"severity": "high", "likelihood": 0.9, "impact": 0.8}, + ) + + assert response["ok"] is True + assert response["priority_score"] > 0.0 + + +def test_llm_intent_analysis_selects_tasks() -> None: + config = EnterprisePipelineConfig( + intent_llm=EnterpriseIntentLLMConfig(provider="custom"), + answer_llm=EnterpriseAnswerLLMConfig(provider="none"), + ) + pipeline = EnterprisePipeline( + query_retriever=DummyRetriever(), + config=config, + intent_llm_client=IntentLLMStub(), + ) + + result = pipeline.review_bundle( + ReviewBundleRequest( + session_id="intent-1", + query_text="Проведи анализ", + requested_tasks=[], + ) + ) + + assert result.executed_tasks == [ReviewTask.RISK_ASSESSMENT, ReviewTask.PROJECT_ERROR_ANALYSIS] + + +def test_answer_llm_can_plan_and_use_tools() -> None: + answer_client = AnswerLLMStub() + config = EnterprisePipelineConfig( + intent_llm=EnterpriseIntentLLMConfig(provider="none"), + answer_llm=EnterpriseAnswerLLMConfig( + provider="custom", + enable_tools=True, + allowed_tools=["score_risks"], + max_tool_calls=1, + ), + ) + pipeline = EnterprisePipeline( + query_retriever=DummyRetriever(), + config=config, + answer_llm_client=answer_client, + ) + + result = pipeline.review_bundle( + ReviewBundleRequest( + session_id="ans-1", + query_text="Сделай вывод и приоритеты", + requested_tasks=[ReviewTask.COMPLIANCE_CHECK], + ) + ) + + assert result.llm_answer == "Generated answer" + assert len(answer_client.prompts) >= 2 + assert "Tool execution results" in answer_client.prompts[-1]