hyperforge-nucliadb 1.0.0.post22__py3-none-any.whl

hyperforge_nucliadb/ask/analysis.py

@@ -0,0 +1,310 @@
+from asyncio import gather
+from copy import deepcopy
+from time import time
+from typing import Any, Dict, List
+
+from hyperforge.manager import Manager
+from hyperforge.memory.memory import QuestionMemory, Source
+from nuclia_models.predict.run_agents import RunTextAgentsRequest
+from nuclia_models.worker.triggers import Relation
+from nucliadb_models import RelationNodeType
+from nucliadb_models.search import Filter, KnowledgeGraphEntity
+
+from hyperforge import PROMPT_ENVIRONMENT
+from hyperforge_nucliadb.ask.config import AskAgentConfig
+from hyperforge_nucliadb.ask.kb_analysis import get_knowledge_base_analysis
+from hyperforge_nucliadb.ask.models import Analysis
+from hyperforge_nucliadb.ask.prompt_analysis import configure_prompts
+from hyperforge_nucliadb.ask.query_analysis import pre_query_analysis
+from hyperforge_nucliadb.driver import NucliaDBDriver
+
+KEYWORD_FILTER_PROMPT = "Extract if any the keywords that should appear on the retrieved results and its a must match, make sure that are keywords that are not common words, and that are not the same as the question or answer. Please define ONLY the keywords without any explanation. Only one or two words maximum. JUST A LIST OF KEYWORDS. DO NOT ADD ANY EXTRA NOTES AT THE END, JUST A LIST OF KEYWORDS"
+
+REPHRASE_SEMANTIC_PROMPT = "Rephrase this question so its better for semantic retrieval, and keep the rephrased question in the same language as the original. Please define ONLY the question without any explanation. JUST A SENTENCE. DO NOT ADD ANY EXTRA NOTES AT THE END, JUST ONE SENTENCE"
+
+REPHRASE_LEXICAL_PROMPT = "Rephrase this question so its better for lexical retrieval, and keep the rephrased question in the same language as the original. Please define ONLY the question without any explanation. JUST A SENTENCE. DO NOT ADD ANY EXTRA NOTES AT THE END, JUST ONE SENTENCE"
+
+
+ASK_JSON_SCHEMA = {
+    "title": "ask_configuration",
+    "description": "Configuration extracted from reasoning engine",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "link": {
+                "type": "boolean",
+                "description": "The user wants link reference to the answer?",
+            },
+            "knowledge_scan": {
+                "type": "string",
+                "description": "If the query requires a knowledge aggregation or scan search to answer define the entities, labels and relations to query in the KB. Example queries: How many ...",
+            },
+            "semantic_query": {
+                "type": "string",
+                "description": REPHRASE_SEMANTIC_PROMPT,
+            },
+            "lexical_query": {"type": "string", "description": REPHRASE_LEXICAL_PROMPT},
+            "visual": {
+                "type": "boolean",
+                "description": "Is required an analysis of an image to answer this question, answer with false or true",
+            },
+            "keywords_filter": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": KEYWORD_FILTER_PROMPT,
+            },
+            "reason": {"type": "string"},
+            "entities": {
+                "type": "array",
+                "description": "Entities related to the user question to query in the KB",
+                "items": {
+                    "type": "string",
+                },
+            },
+            "relations": {
+                "type": "array",
+                "description": "Relations related to the user question to query in the KB",
+                "items": {
+                    "type": "string",
+                },
+            },
+            "pre_queries": {
+                "type": "array",
+                "items": {
+                    "type": "string",
+                },
+                "description": "Pre queries to run before the main query to gather more information",
+            },
+        },
+        "required": [
+            "semantic_query",
+            "lexical_query",
+            "visual",
+            "keywords_filter",
+            "reason",
+            "pre_queries",
+        ],
+    },
+}
+
+
+ASK_PROMPT = """
+
+Information about the KB:
+
+# {{sid}}
+
+description: {{sdescription}}
+
+And given the question: {{question}}
+
+## labels: {{slabels}}
+
+## Facets:
+
+{{facets}}
+
+# Important rules to follow
+
+{% for rule in rules %}
+{{rule}}
+{% endfor -%}
+
+
+"""
+
+ASK_PROMPT_TEMPLATE = PROMPT_ENVIRONMENT.from_string(ASK_PROMPT)
+
+
+async def question_analysis(
+    memory: QuestionMemory,
+    manager: Manager,
+    nucliadb_driver: NucliaDBDriver,
+    source: Source,
+    config: AskAgentConfig,
+    question: str,
+    step_title: str = "Knowledge Box Ask: Choose parameters",
+):
+    if question is None:
+        raise Exception("Question is None")
+    question = question.strip()
+
+    if source is None:
+        raise Exception("Source is None")
+
+    kb_analysis = get_knowledge_base_analysis(source)
+
+    query_analysis_result = pre_query_analysis(config, question, kb_analysis)
+
+    ask_json_schema_copy: Dict[str, Any] = deepcopy(ASK_JSON_SCHEMA)
+    configure_prompts(config, ask_json_schema_copy, kb_analysis, query_analysis_result)
+
+    facets = ""
+    facets += "## Content Types\n"
+    facets += "The following content types are available in the KB:\n\n"
+    for key, count in kb_analysis.content_types.items():
+        facets += f"- {key}: {count}\n"
+
+    facets += "The following languages are available in the KB:\n\n"
+    for key, count in kb_analysis.languages.items():
+        facets += f"- {key}: {count}\n"
+
+    prompt = ASK_PROMPT_TEMPLATE.render(
+        source=source,
+        rules=memory.get_rules(),
+        question=question,
+        facets=facets,
+        sid=source.id,
+        sdescription=source.description,
+        slabels=source.labels,
+    )
+
+    t0 = time()
+    input_tokens = 0.0
+    output_tokens = 0.0
+
+    task1 = manager.execute_json(
+        user_id="arag-ask",
+        prompt=prompt,
+        model=config.configuration_model,
+        schema=ask_json_schema_copy,
+        tracking=memory.get_tracking_info(),
+    )
+
+    task2 = nucliadb_driver.driver.run_text_agents(
+        kbid=nucliadb_driver.config.kbid,
+        content=RunTextAgentsRequest(user_id="arag-ask", texts=[question]),
+    )
+
+    task3 = manager.tokens_predict(
+        text=question,
+        model=kb_analysis.entity_model,
+    )
+
+    # Extract configuration from reasoning engine
+    (configuration_json, input, output), das_answer, tokens = await gather(
+        task1, task2, task3
+    )
+
+    input_tokens += input
+    output_tokens += output
+
+    entities_to_filter = deepcopy(config.query_entities)
+    relations_to_filter: List[Relation] = []
+    labels_to_filter: List[Filter] = []
+    if len(config.filters) > 0:
+        labels_to_filter.append(Filter(all=config.filters))
+    new_text_fields: List[str] = []
+    new_json_fields: Dict[str, Any] = {}
+
+    for detected_token in tokens.tokens:
+        entities_to_filter.append(
+            KnowledgeGraphEntity(
+                name=detected_token.text,
+                type=RelationNodeType.ENTITY,
+                subtype=detected_token.ner,
+            )
+        )
+
+    for result in das_answer.results:
+        input_tokens += result.input_nuclia_tokens
+        output_tokens += result.output_nuclia_tokens
+        for payload in result.payloads:
+            # we could implement for guards
+            if payload.entities is not None:
+                for entity in payload.entities:
+                    for entity_class, entity_value in entity.labels.items():
+                        entities_to_filter.append(
+                            KnowledgeGraphEntity(
+                                name=entity_value,
+                                type=RelationNodeType.ENTITY,
+                                subtype=entity_class,
+                            )
+                        )
+
+            if payload.relations is not None:
+                for relations in payload.relations:
+                    relations_to_filter.extend(relations.relations)
+
+            if payload.labels is not None:
+                das_filters: list[str] = []
+                for label in payload.labels:
+                    if isinstance(label.labels, list):
+                        for label_obj_str in label.labels:
+                            das_filters.append(label_obj_str)
+                    elif isinstance(label.labels, dict):
+                        for key, value in label.labels.items():
+                            das_filters.append(f"{key}:{value}")
+                    else:
+                        raise Exception(f"Unknown label type: {label.labels}")
+                labels_to_filter.append(Filter(all=das_filters))
+
+            if payload.asks is not None:
+                for ask in payload.asks:
+                    if ask.empty is False and ask.text is not None:
+                        new_text_fields.append(ask.text)
+                    elif ask.empty is False and ask.json_output is not None:
+                        new_json_fields.update(ask.json_output)
+
+    for entity_llm in configuration_json.get("entities", []):
+        entities_to_filter.append(
+            KnowledgeGraphEntity(
+                name=entity_llm,
+                type=RelationNodeType.ENTITY,
+            )
+        )
+
+    if kb_analysis.default_semantic_model is None and (
+        kb_analysis.semantic_models is None or len(kb_analysis.semantic_models) == 0
+    ):
+        raise Exception("No semantic model found in the knowledge base")
+
+    semantic_model = (
+        kb_analysis.default_semantic_model
+        if kb_analysis.default_semantic_model
+        else kb_analysis.semantic_models[0]
+    )
+
+    semantic_query = configuration_json["semantic_query"]
+    lexical_query = configuration_json["lexical_query"]
+    visual = configuration_json.get("visual", False)
+    link = configuration_json.get("link", False)
+    knowledge_graph = configuration_json.get("knowledge_graph", None)
+    keywords_filter: List[str] = configuration_json.get("keywords_filter", [])
+    if isinstance(keywords_filter, str):
+        keywords_filter = [word.strip() for word in keywords_filter.split(",")]
+
+    reason = configuration_json.get("reason")
+    pre_queries = configuration_json.get("pre_queries", [])
+    relations = configuration_json.get("relations", [])
+    await memory.add_step(
+        step_module="ask",
+        step_title=step_title,
+        step_agent_path=f"/context/{config.id if config.id else 'default'}",
+        step_value=str(configuration_json),
+        step_reason=reason if reason is not None else "",
+        timeit=time() - t0,
+        input_nuclia_tokens=input_tokens,
+        output_nuclia_tokens=output_tokens,
+    )
+
+    return Analysis(
+        semantic_query=semantic_query,
+        lexical_query=lexical_query,
+        visual=visual,
+        keywords_filter=keywords_filter,
+        reason=reason,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        labels=labels_to_filter,
+        semantic_model=semantic_model,
+        entities=entities_to_filter,
+        pre_queries=pre_queries,
+        relations=relations_to_filter,
+        new_text_fields=new_text_fields,
+        new_json_fields=new_json_fields,
+        semantic_config=kb_analysis.semantic_model_configs[semantic_model],
+        top_k=config.top_k,
+        link=link,
+        knowledge_graph=knowledge_graph,
+    )

hyperforge_nucliadb/ask/ask.py

@@ -0,0 +1,326 @@
+import asyncio
+from time import time
+from typing import Any, ClassVar, Dict, List, Optional, cast
+
+from hyperforge.agent import Agent
+from hyperforge.configure import agent
+from hyperforge.context.agent import ContextAgent, trace_agent
+from hyperforge.definition import FunctionDefinition
+from hyperforge.manager import Manager
+from hyperforge.memory.memory import Context, QuestionMemory, Source
+
+from hyperforge import logger
+from hyperforge_nucliadb.ask.analysis import question_analysis
+from hyperforge_nucliadb.ask.config import AskAgentConfig
+from hyperforge_nucliadb.ask.hydrate import hydrate
+from hyperforge_nucliadb.ask.knowledge_scan import knowledge_scan
+from hyperforge_nucliadb.ask.models import Analysis
+from hyperforge_nucliadb.ask.multi import choose_source
+from hyperforge_nucliadb.ask.nucliadb import query_ndb, standard_query_ndb
+from hyperforge_nucliadb.ask.rerank import rerank
+from hyperforge_nucliadb.driver import NucliaDBDriver
+
+
+@agent(
+    id="ask",
+    agent_type="context",
+    title="Knowledge Box Ask",
+    description="Ask a question to the knowledge box and retrieve relevant information",
+    config_schema=AskAgentConfig,
+)
+class AskAgent(ContextAgent, Agent[AskAgentConfig]):
+    __published_functions__: ClassVar[Dict[str, FunctionDefinition]] = {
+        "search_by_title": FunctionDefinition(
+            name="search_by_title",
+            description="Search for context in the Knowledge Box by title. Useful for specific queries where the title is known.",
+            parameters={
+                "title": {
+                    "type": "string",
+                    "description": "The title to search for in the Knowledge Box.",
+                },
+                "filters": {
+                    "type": "array",
+                    "items": {"type": "string"},
+                    "description": "Filters to apply when searching in the Knowledge Box.",
+                },
+            },
+        ),
+        "ask_analysis_query": FunctionDefinition(
+            name="ask_analysis_query",
+            description="Generate filters to be used when querying the Knowledge Box. Useful to refine search results based on the question.",
+            parameters={
+                "question": {
+                    "type": "string",
+                    "description": "The question based on which to generate filters.",
+                },
+                "kbid": {
+                    "type": "string",
+                    "description": "The Knowledge Box ID to use for generating the analysis. If not provided, the default sources will be used.",
+                },
+            },
+        ),
+        "ask_agent": FunctionDefinition(
+            name="ask_agent",
+            description="Search for context in the Knowledge Box by title. Useful for specific queries where the title is known.",
+            parameters={
+                "question": {
+                    "type": "string",
+                    "description": "The question to search for in the Knowledge Box.",
+                },
+            },
+        ),
+    }
+
+    fallback: Optional[ContextAgent] = None
+
+    async def ask_analysis_query(
+        self,
+        memory: QuestionMemory,
+        manager: Manager,
+        question: str,
+        **kwargs: Any,
+    ) -> List[Analysis]:
+        sources = await self.get_source(question, memory, manager)
+        result = []
+        for source in sources:
+            nucliadb_driver: Optional[NucliaDBDriver] = cast(
+                NucliaDBDriver, manager.drivers.get(source.id)
+            )
+            if nucliadb_driver is None:
+                raise Exception("No NDB available")
+            analysis = await question_analysis(
+                memory,
+                manager,
+                nucliadb_driver,
+                source,
+                config=self.config,
+                question=question,
+                step_title=self.step_title("Choose parameters"),
+            )  # type: ignore
+            result.append(analysis)
+        return result
+
+    async def work(
+        self,
+        memory: QuestionMemory,
+        manager: Manager,
+        context: Context,
+        nucliadb_driver: NucliaDBDriver,
+        source: Source,
+        question: str,
+    ):
+        # Generate semantic questions required to answer
+        analysis: Optional[Analysis] = None
+        if self.config.ai_parameter_search:
+            try:
+                analysis = await question_analysis(
+                    memory,
+                    manager,
+                    nucliadb_driver,
+                    source,
+                    config=self.config,
+                    question=question,
+                    step_title=self.step_title("Choose parameters"),
+                )
+            except Exception as e:
+                logger.exception("Error analyzing question")
+                raise e
+
+            try:
+                if analysis.knowledge_graph is not None:
+                    search_result = await knowledge_scan(
+                        memory,
+                        manager,
+                        analysis,
+                        nucliadb_driver,
+                        config=self.config,
+                    )
+                else:
+                    search_result = await query_ndb(
+                        memory,
+                        manager,
+                        analysis,
+                        nucliadb_driver,
+                        config=self.config,
+                    )
+
+            except Exception as e:
+                logger.exception("Error querying NDB")
+                raise e
+        else:
+            search_result = await standard_query_ndb(
+                memory,
+                manager,
+                nucliadb_driver,
+                config=self.config,
+                question=question,
+            )
+
+        await hydrate(
+            context,
+            search_result,
+            nucliadb_driver=nucliadb_driver,
+            vllm=self.config.vllm,
+            after=self.config.after,
+            before=self.config.before,
+            full_resource=self.config.full_resource,
+            visual=analysis.visual if analysis is not None else False,
+            link=analysis.link if analysis is not None else False,
+            module=self.config.module,
+        )
+        await rerank(context, manager=manager, top_k=analysis.top_k if analysis else 20)
+
+    async def get_source(
+        self, question: str, memory: QuestionMemory, manager: Manager
+    ) -> List[Source]:
+        if len(self.config.sources) == 0:
+            # Do we really want to do this without telling the user?
+            sources = manager.nucliadbs()
+        else:
+            sources = self.config.sources
+
+        # Only makes sense if we have more than one source
+        chosen_sources = await choose_source(
+            memory,
+            manager,
+            sources,
+            question,
+            ident=self.config.id if self.config.id else "default",
+            step_title=self.step_title("Choose sources"),
+        )
+        return chosen_sources
+
+    @trace_agent
+    async def get_question_context(
+        self,
+        memory: QuestionMemory,
+        manager: Manager,
+        question_uuid: str,
+        question: str,
+        flow_id: str,
+        extra_context: Optional[Dict[str, Any]] = None,
+    ):
+        if len(self.config.sources) == 0:
+            # Do we really want to do this without telling the user?
+            sources = manager.nucliadbs()
+        else:
+            sources = self.config.sources
+
+        # Only makes sense if we have more than one source
+        chosen_sources = await choose_source(
+            memory,
+            manager,
+            sources,
+            question,
+            ident=self.config.id if self.config.id else "default",
+            step_title=self.step_title("Choose sources"),
+        )
+
+        await asyncio.gather(
+            *[
+                self.rag(
+                    memory, manager, source, question_uuid, question, flow_id=flow_id
+                )
+                for source in chosen_sources
+            ]
+        )
+
+    async def rag(
+        self,
+        memory: QuestionMemory,
+        manager: Manager,
+        source_obj: Source,
+        question_uuid: str,
+        question: str,
+        flow_id: str,
+    ):
+        source = source_obj.id
+        nucliadb_driver: Optional[NucliaDBDriver] = cast(
+            NucliaDBDriver, manager.drivers.get(source)
+        )
+        if nucliadb_driver is None:
+            raise Exception("No NDB available")
+
+        context = Context(
+            agent_id=self.config.id,
+            original_question_uuid=memory.original_question_uuid,
+            actual_question_uuid=question_uuid,
+            question=question,
+            source=source,
+            agent="ask",
+            title=self.config.title
+            if self.config.title
+            else f"Retrieval on {source} Knowledge Box",
+        )
+
+        await self.work(memory, manager, context, nucliadb_driver, source_obj, question)
+        ident = self.config.id if self.config.id else "default"
+        # If no chunks were found, we can skip the summarization
+        if context.chunks is not None and len(context.chunks) != 0:
+            # These missing questions are not used as of now
+            missing_question = await self.save_ctx_and_return_missing(
+                context=context,
+                question=question,
+                memory=memory,
+                manager=manager,
+                flow_id=flow_id,
+            )
+        else:
+            missing_question = (question_uuid, question)
+            await memory.add_step(
+                step_module=self.config.module,
+                step_title=self.step_title("NucliaDB context"),
+                step_reason="No context found",
+                step_value="Unsuccessful retrieval",
+                timeit=time() - time(),
+                step_agent_path=f"/context/{ident}",
+            )
+        # Query NDB
+        if missing_question is not None:
+            missing_uuid, missing = missing_question
+            # This is always failing back to another nucliadb call in case of missing information
+            # TODO: decide if we want to do this or not
+            # implicit fallback to another NDB call may introduce unnecessary delay
+            # for now I'm only adding it if we have no fallback agent configured
+            if self.config.fallback is None:
+                context = Context(
+                    agent_id=self.config.id,
+                    original_question_uuid=memory.original_question_uuid,
+                    actual_question_uuid=missing_uuid,
+                    question=missing,
+                    source=source,
+                    agent="ask",
+                    title=self.config.title
+                    if self.config.title
+                    else f"Retrieval on {source} Knowledge Box",
+                )
+
+                await self.work(
+                    memory,
+                    manager,
+                    context,
+                    nucliadb_driver,
+                    source_obj,
+                    question,
+                )
+                new_missing_questions = await self.save_ctx_and_return_missing(
+                    context=context,
+                    question=missing,
+                    memory=memory,
+                    manager=manager,
+                    flow_id=flow_id,
+                )
+                if new_missing_questions is not None:
+                    # TODO: add proper structure to notify the user that the answer is missing
+                    logger.warning(
+                        f"Missing information after fallback: {[m[1] for m in new_missing_questions]}. No fallback agent configured."
+                    )
+            elif self.fallback is not None:
+                # TODO: Add next agent
+                await self.fallback.get_question_context(
+                    memory,
+                    manager,
+                    question=missing,
+                    question_uuid=missing_uuid,
+                    flow_id=flow_id,
+                )

hyperforge_nucliadb/ask/config.py

@@ -0,0 +1,52 @@
+from typing import List, Literal, Optional
+
+from hyperforge.context.config import ContextAgentConfig
+from hyperforge.utils import WidgetType
+from nucliadb_models.search import KnowledgeGraphEntity
+from pydantic import Field
+from pydantic.config import ConfigDict
+
+
+class AskAgentConfig(ContextAgentConfig):
+    model_config = ConfigDict(title="Knowledge Box Ask")
+    module: Literal["ask"] = "ask"
+    published_functions: Optional[tuple[str, ...]] = Field(
+        default=("search_by_title", "ask_analysis_query", "ask_agent"),
+        title="Published functions",
+        description="List of functions published by this agent to be used by other agents in the chain",
+        json_schema_extra={
+            "widget": WidgetType.NOT_SHOWN,
+        },
+    )
+    pre_queries: List[str] = Field(default_factory=list)
+    filters: List[str] = Field(default_factory=list)
+    security_groups: List[str] = Field(default_factory=list)
+    rephrase_semantic_custom_prompt: Optional[str] = None
+    rephrase_lexical_custom_prompt: Optional[str] = None
+    keywords_custom_prompt: Optional[str] = None
+    visual_enable_prompt: Optional[str] = None
+    date_range_enabled: bool = False
+    before: int = 2
+    after: int = 2
+    top_k: int = 20
+    extra_fields: List[str] = Field(default_factory=list)
+    full_resource: bool = False
+    vllm: bool = True
+    query_entities: List[KnowledgeGraphEntity] = Field(default_factory=list)
+    retrieve_related: Optional[str] = None
+    sources: List[str] = Field(
+        default_factory=list,
+        json_schema_extra={
+            "show_in_node": True,
+        },
+    )
+    configuration_model: str = Field(
+        default="gemini-2.5-flash",
+        title="Generative model",
+        description="Model used to generate the configuration",
+        json_schema_extra={"widget": WidgetType.MODEL_SELECT},
+    )
+    fast_answer: bool = True
+    # Setting this to True so that this PR will not break the current behavior
+    # of the agent
+    ai_parameter_search: bool = True