hyperforge 1.0.0.post19__py3-none-any.whl

hyperforge/context/agent.py

@@ -0,0 +1,377 @@
+import string
+from functools import lru_cache
+from time import time
+from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, cast
+from uuid import uuid4
+
+from hyperforge import logger
+from hyperforge.configure import get_agent_klass
+from hyperforge.context.config import ContextAgentConfig
+from hyperforge.definition import FunctionDefinition
+from hyperforge.manager import Manager
+from hyperforge.memory import Context, QuestionMemory
+from hyperforge.prompts import (
+    NEXT_REPHRASE_JSON_SCHEMA,
+    NEXT_REPHRASE_PROMPT_SYSTEM,
+    NEXT_REPHRASE_PROMPT_TEMPLATE,
+    VALIDATE_JSON_SCHEMA,
+    VALIDATE_OR_ANSWER_PROMPT_TEMPLATE,
+)
+from hyperforge.trace import trace_agent
+
+
+@lru_cache(maxsize=64)
+def generate_ctx_block_id(n: int) -> str:
+    """Generates block ids in the form block-AA, block-AB, ..., block-AZ, block-BA, ..., block-ZZ"""
+    # XXX: We use letters to identify block since uuids confuse the LLMs and with numbers it confuses the block numbers with footnote numbers
+    if n < 0:
+        raise ValueError("Number must be non-negative")
+    if n >= 26 * 26:
+        logger.warning(
+            "Block ID exceeds maximum limit for citations,", extra={"block_id": n}
+        )
+        n = n % (26 * 26)
+    letters = string.ascii_uppercase
+    first = letters[(n // 26)]
+    second = letters[n % 26]
+    return "block-" + first + second
+
+
+async def build_context_agent(config: ContextAgentConfig) -> "ContextAgent":
+    agent_class = get_agent_klass(config.module)
+    assert issubclass(agent_class, ContextAgent), (
+        f"Agent {config.module} is not a ContextAgent"
+    )
+    return await agent_class.from_config(config)
+
+
+class ContextAgent:
+    fallback: Optional["ContextAgent"] = None
+    next_agent: Optional["ContextAgent"] = None
+    agent_description: str = "Agent that provides context to answer questions."
+    __published_functions__: ClassVar[Dict[str, FunctionDefinition]] = {}
+    exposed_functions: Optional[List[str]] = None
+    agent_id: str
+
+    async def preload(self, manager: "Manager", memory: "QuestionMemory") -> None:
+        """Lifecycle hook called by SmartAgent before tool discovery.
+
+        Override in subclasses that need a live connection (e.g. MCP) to
+        populate ``__published_functions__`` at runtime.  The default
+        implementation is a no-op.
+        """
+        pass
+
+    @property
+    def context_config(self) -> ContextAgentConfig:
+        return cast(ContextAgentConfig, self.config)  # type: ignore
+
+    async def inner_from_config(self, config: Any, agent_id: Optional[str] = None):
+        await self.context_from_config(config)
+
+    async def context_from_config(self, config: ContextAgentConfig):
+        # Shared logic for fallback and next_agent
+        if config.fallback:
+            self.fallback = await build_context_agent(config.fallback)
+        if config.next_agent:
+            self.next_agent = await build_context_agent(config.next_agent)
+
+    async def validate_ctx_and_answer(
+        self,
+        memory: QuestionMemory,
+        manager: Manager,
+        context: Context,
+        question: str,
+        images: bool = False,
+        fast_answer: bool = False,
+        user_id: str = "rao_answer_summary",
+        use_stored_context_prompts: bool = False,
+        title: Optional[str] = None,
+    ) -> Tuple[Literal["yes", "no", "error"], Optional[str], Optional[str]]:
+        """
+        Validate if the context is useful to answer the question and attempt use it to answer the user's question.
+
+        Reports back if the context is useful, what is missing, any error during summarization.
+
+        It will also save the answer attempt and the citations in the context object.
+        """
+        t0 = time()
+        model = self.context_config.context_validation_model
+        module = self.context_config.module
+        ident = self.context_config.id if self.context_config.id else "default"
+
+        citations = None
+        contexts = [x.render() for x in context.chunks]
+        contexts.extend([x for x in context.structured])
+
+        contexts_items = {generate_ctx_block_id(i): x for i, x in enumerate(contexts)}
+        contexts_items.update(
+            {
+                generate_ctx_block_id(i + len(context.chunks)): x
+                for i, x in enumerate(context.structured)
+            }
+        )
+        block_id_to_chunk = {
+            generate_ctx_block_id(i): chunk.chunk_id
+            for i, chunk in enumerate(context.chunks)
+        }
+        block_id_to_chunk.update(
+            {
+                generate_ctx_block_id(i + len(context.chunks)): f"structured-{i}"
+                for i in range(len(context.structured))
+            }
+        )
+
+        extra_prompts = [prompt.render() for prompt in context.prompts]
+        prompt = VALIDATE_OR_ANSWER_PROMPT_TEMPLATE.render(
+            question=question,
+            contexts=contexts_items,
+            extra_prompts=extra_prompts if use_stored_context_prompts else "",
+        )
+        try:
+            (
+                response,
+                input_tokens,
+                output_tokens,
+            ) = await manager.execute_json(
+                user_id=user_id + f"-{module}",
+                images=context.images if images else {},
+                prompt=prompt,
+                schema=VALIDATE_JSON_SCHEMA,
+                model=model,
+            )
+            answer = response.get("answer", "")
+            missing = response.get("missing_info_query", None)
+            useful = response.get("useful", "yes")
+            reason = response.get("reason", "")
+            citations = response.get("citations", None)
+        except Exception as e:
+            error_message = f"Error executing validation of contexts for {module} with question '{question}': {e}"
+            await memory.add_step(
+                step_module=module,
+                step_title=f"{title}: Context summary error",
+                step_reason="Error in summary generation",
+                step_value="Error in summary generation",
+                timeit=time() - t0,
+                step_agent_path=f"/context/{ident}",
+                error=error_message,
+            )
+            return "error", question, error_message
+
+        # removing this for now
+        # if useful == "yes" and fast_answer:
+        #     await memory.add_answer(answer, module="ask", agent_path=f"/context/{ident}")
+
+        if citations is not None:
+            cited_chunks = [
+                block_id_to_chunk[citation]
+                for citation in set(citations)
+                if citation in block_id_to_chunk
+            ]
+            context.citations = cited_chunks
+
+        context.summary = (
+            answer if "not enough data to answer this" not in answer.lower() else ""
+        )
+
+        await memory.add_step(
+            step_module=module,
+            step_title=f"{title}: Context validation",
+            step_reason=reason,
+            step_value="No missing information" if not missing else missing,
+            timeit=time() - t0,
+            input_nuclia_tokens=input_tokens,
+            output_nuclia_tokens=output_tokens,
+            step_agent_path=f"/context/{ident}",
+        )
+
+        if missing is not None and missing.strip() != "":
+            return useful, missing, None
+        return useful, None, None
+
+    async def rephrase(
+        self,
+        memory: QuestionMemory,
+        manager: Manager,
+        contexts: Dict[str, Any],
+        ident: str,
+        question: str,
+        question_uuid: str,
+        model: str = "",
+        module: str = "agent",
+        user_id: str = "next",
+        title: Optional[str] = None,
+    ) -> Tuple[str, str]:
+        t0 = time()
+        rephrased = False
+        prompt = NEXT_REPHRASE_PROMPT_TEMPLATE.render(
+            question=question,
+            contexts=list(contexts.values()),
+            info=self.agent_description,
+            extra_info=self.context_config.context_aware_rephrasing_prompt,
+        )
+        try:
+            (
+                response,
+                input_tokens,
+                output_tokens,
+            ) = await manager.execute_json(
+                user_id=user_id + f"-{module}",
+                system=NEXT_REPHRASE_PROMPT_SYSTEM,
+                prompt=prompt,
+                schema=NEXT_REPHRASE_JSON_SCHEMA,
+                model=model,
+            )
+            rephrased_question = response.get("rephrased_question", "")
+            needed = response.get("needed", True)
+            reason = response.get("reason", "")
+        except Exception as e:
+            error_message = f"Error executing rephrase when executing next agent {module} with question '{question}': {e}"
+            await memory.add_step(
+                step_module=module,
+                step_title=f"{title}: Rephrase error",
+                step_reason="Error in rephrase generation",
+                step_value="Error in rephrase generation",
+                timeit=time() - t0,
+                step_agent_path=f"/context/{ident}",
+                error=error_message,
+            )
+            return question, question_uuid
+        if (
+            rephrased_question.strip() != ""
+            and needed is True
+            and "not enough data to answer this" not in rephrased_question.lower()
+        ):
+            rephrased_question_uuid = uuid4().hex
+            rephrased = True
+
+        await memory.add_step(
+            step_module=module,
+            step_title=f"{title}: Rephrase",
+            step_reason=reason,
+            step_value="No need for rephrasing"
+            if rephrased is False
+            else rephrased_question,
+            timeit=time() - t0,
+            input_nuclia_tokens=input_tokens,
+            output_nuclia_tokens=output_tokens,
+            step_agent_path=f"/context/{ident}",
+        )
+        return (
+            rephrased_question if rephrased else question,
+            rephrased_question_uuid if rephrased else question_uuid,
+        )
+
+    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,
+    ) -> List[tuple[str, str]]:
+        """To be implemented by child classes to get context for the question.
+        Should return a list of (question_uuid, question) tuples representing missing questions.
+        """
+        raise NotImplementedError()
+
+    @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 extra_context is not None:
+            question, question_uuid = await self.rephrase(
+                memory=memory,
+                manager=manager,
+                question_uuid=question_uuid,
+                question=question,
+                contexts=extra_context,
+                model=self.context_config.rephrase_model,
+                module=self.context_config.module,
+                user_id="next_rephrase",
+                ident=self.agent_id,
+            )
+
+        # Missing questions should be a list of (question_uuid, question) tuples
+        missing_questions = await self._get_question_context(
+            memory, manager, question_uuid, question, flow_id=flow_id
+        )
+
+        if self.fallback is not None:
+            for missing_uuid, missing in missing_questions:
+                await self.fallback.get_question_context(
+                    memory,
+                    manager,
+                    question_uuid=missing_uuid,
+                    question=missing,
+                    flow_id=flow_id,
+                )
+
+        if self.next_agent is not None:
+            extra_context = extra_context or {}
+            for agent in [self, self.fallback]:
+                if agent:
+                    answer_summaries = memory.get_agent_answer_summaries(
+                        flow_id=flow_id, agent_id=agent.agent_id
+                    )
+                    if answer_summaries:
+                        extra_context[agent.agent_id] = "\n".join(answer_summaries)
+
+            await self.next_agent.get_question_context(
+                memory,
+                manager,
+                question_uuid,
+                question,
+                extra_context=extra_context,
+                flow_id=flow_id,
+            )
+
+    async def save_ctx_and_return_missing(
+        self,
+        *,
+        memory: QuestionMemory,
+        manager: Manager,
+        question: str,
+        context: Context,
+        flow_id: str,
+        use_stored_context_prompts: bool = False,
+    ) -> tuple[str, str] | None:
+        # We trigger context validation and answer attempt only if we have a fallback or next agent or pruning is enabled
+        if (
+            self.fallback is not None
+            or self.next_agent is not None
+            or self.context_config.prune_context
+            or use_stored_context_prompts
+        ) and self.context_config.context_validation_model:
+            (
+                useful,
+                missing,
+                validate_error,
+            ) = await self.validate_ctx_and_answer(
+                memory,
+                manager,
+                context,
+                question=question,
+                use_stored_context_prompts=use_stored_context_prompts,
+            )
+
+            if useful == "yes" or validate_error is not None:
+                if self.context_config.prune_context and validate_error is None:
+                    context.prune_to_citations()
+                await memory.save_context(flow_id=flow_id, context=context)
+            missing_question = (
+                None
+                if missing is None or missing.strip() == ""
+                else (uuid4().hex, missing)
+            )
+            return missing_question
+        else:
+            await memory.save_context(flow_id=flow_id, context=context)
+        return None

hyperforge/context/config.py

@@ -0,0 +1,103 @@
+from typing import Any, Dict, Optional, Tuple
+
+from pydantic import (
+    AliasChoices,
+    BaseModel,
+    Field,
+    ValidationError,
+    ValidatorFunctionWrapHandler,
+    field_serializer,
+    field_validator,
+)
+
+from hyperforge.agent import AgentConfig
+from hyperforge.configure import get_agent_config_klass
+from hyperforge.utils import WidgetType
+
+
+def context_agent_validator(value: Any, handler: ValidatorFunctionWrapHandler) -> str:
+    try:
+        return handler(value)
+    except ValidationError as err:
+        if err.errors()[0]["type"] == "string_too_long":
+            return handler(value[:5])
+        else:
+            raise
+
+
+class ContextAgentConfig(AgentConfig):
+    fallback: Optional["ContextAgentConfig"] = Field(
+        default=None,
+        title="Fallback agent",
+        description="Agent to use in case this one fails",
+        json_schema_extra={
+            "widget": WidgetType.NOT_SHOWN,
+        },
+    )
+    next_agent: Optional["ContextAgentConfig"] = Field(
+        default=None,
+        title="Next agent",
+        description="Agent to run after executing the code",
+        json_schema_extra={
+            "widget": WidgetType.NOT_SHOWN,
+        },
+    )
+    context_validation_model: str = Field(
+        default="chatgpt-azure-4o-mini",
+        title="Context validation model",
+        description="Model used to validate the agent's generated context and generate an answer attempt to the user's question.",
+        json_schema_extra={"widget": WidgetType.MODEL_SELECT},
+        # Backward compatibility with frontend and stored configurations
+        validation_alias=AliasChoices("context_validation_model", "summarize_model"),
+    )
+    rephrase_model: str = Field(
+        default="chatgpt-azure-4o-mini",
+        title="Rephrase model",
+        description="Model used to rephrase the question based on context (only used in agents that follow others in a chain)",
+        json_schema_extra={"widget": WidgetType.MODEL_SELECT},
+    )
+    context_aware_rephrasing_prompt: Optional[str] = Field(
+        default=None,
+        title="Context-aware rephrasing prompt",
+        description="Custom prompt to give the model when rephrasing questions based on previous contexts (only used in agents that follow others in a chain)",
+        json_schema_extra={
+            "widget": WidgetType.EXPANDABLE_TEXTAREA,
+        },
+    )
+    prune_context: bool = Field(
+        default=True,
+        title="Prune context",
+        description="Whether to prune the context pieces generated by this agent before storing them in memory. Will always trigger a context validation call if enabled.",
+    )
+
+    published_functions: Optional[Tuple[str, ...]] = Field(
+        default=None,
+        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,
+        },
+    )
+    # always_validate_contexts: bool = Field(
+    #     default=False,
+    #     title="Always validate and attempt to generate an answer",
+    #     description="Whether to always validate the context retrieved and attempt to generate an answer to the question with this agent's context. Regardless if it is not required for pruning or for the next agent in the chain.",
+    # )
+
+    @field_validator("fallback", "next_agent", mode="before")
+    @classmethod
+    def is_context_agent(cls, value: Dict[str, Any]) -> BaseModel:
+        if value is None:
+            return value
+        module = value.get("module")
+        if module is None:
+            raise ValueError("Invalid agent config: missing 'module' field")
+
+        agent_config_klass = get_agent_config_klass(module, agent_type="context")
+        return agent_config_klass.model_validate(value)
+
+    @field_serializer("fallback", "next_agent")
+    def serialize_context_agent(self, field: BaseModel) -> Optional[Dict[str, Any]]:
+        if field is None:
+            return field
+        return field.model_dump()

hyperforge/database.py

@@ -0,0 +1,3 @@
+import sqlalchemy as sa
+
+metadata = sa.MetaData()

hyperforge/db/__init__.py

@@ -0,0 +1,6 @@
+import logging
+
+logger = logging.getLogger("hyperforge.db")
+
+
+CONFIG_CACHE_KEY = "wmc_{account}_{kbid}_{task_type}"