ws-bom-robot-app 0.0.37__py3-none-any.whl → 0.0.103__py3-none-any.whl

ws_bom_robot_app/llm/agent_handler.py

@@ -1,180 +1,176 @@
-from asyncio import Queue
-from langchain_core.agents import AgentFinish
-from langchain_core.outputs import ChatGenerationChunk, GenerationChunk
-from langchain.callbacks.base import AsyncCallbackHandler
-from ws_bom_robot_app.llm.utils.print import printJson, printString
-from typing import Any, Dict, List, Optional, Union
-from uuid import UUID
-import ws_bom_robot_app.llm.settings as settings
-from langchain_core.callbacks.base import AsyncCallbackHandler
-from langchain_core.outputs import ChatGenerationChunk, GenerationChunk
-from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
-import json
-
-# Here is a custom handler that will print the tokens to stdout.
-# Instead of printing to stdout you can send the data elsewhere; e.g., to a streaming API response
-
-
-def _parse_token(llm:str,token: str) -> str:
-    """Parses the token based on the LLM provider."""
-    if llm == "anthropic" and isinstance(token, list):
-      first = token[0]
-      if 'text' in first:
-        token = first['text']
-      else:
-        #[{'id': 'toolu_01GGLwJcrQ8PvFMUkQPGu8n7', 'input': {}, 'name': 'document_retriever_xxx', 'type': 'tool_use', 'index': 1}]
-        token = ""
-    return token
-
-class AgentHandler(AsyncCallbackHandler):
-
-    def __init__(self, queue: Queue, llm:str, threadId: str = None) -> None:
-        super().__init__()
-        self._threadId = threadId
-        self.json_block = ""
-        self.is_json_block = False
-        self.backtick_count = 0  # Conteggio dei backticks per il controllo accurato
-        self.queue = queue
-        self.llm = llm
-
-    async def on_llm_start(
-        self,
-        serialized: Dict[str, Any],
-        prompts: List[str],
-        *,
-        run_id: UUID,
-        parent_run_id: UUID = None,
-        tags: List[str] = None,
-        metadata: Dict[str, Any] = None,
-        **kwargs: Any,
-    ) -> None:
-        firstChunk = {
-            "type": "info",
-            "threadId": self._threadId,
-        }
-        await self.queue.put(printString(firstChunk))
-
-    """async def on_chat_model_start(self, serialized: Dict[str, Any], messages: List[List[BaseMessage]], *, run_id: UUID = None, parent_run_id = None, tags = None, metadata = None, **kwargs: Any) -> Any:
-        pass"""
-
-    async def on_tool_end(self, output: Any, *, run_id: UUID, parent_run_id: UUID = None, tags: List[str] = None, **kwargs: Any) -> None:
-      pass
-
-    async def on_llm_new_token(
-        self,
-        token: str,
-        *,
-        chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]] = None,
-        run_id: UUID,
-        parent_run_id: Optional[UUID] = None,
-        tags: Optional[List[str]] = None,
-        **kwargs: Any,
-    ) -> None:
-        """Gestisce i nuovi token durante lo streaming."""
-
-        if token:
-            token = _parse_token(self.llm,token)
-            if token:
-              self.backtick_count += token.count("`")
-
-              if self.backtick_count >= 3:
-                  if not self.is_json_block:
-                      self.is_json_block = True
-                      self.json_block = ""
-                  else:
-                      self.is_json_block = False
-                      self.json_block += token.replace("```json", '')
-                      await self.process_json_block(self.json_block)
-                      self.json_block = ""
-                  self.backtick_count = 0
-              elif self.is_json_block:
-                  self.json_block += token
-              else:
-                  await self.queue.put(printString(token))
-
-    async def on_agent_finish(
-        self,
-        finish: AgentFinish,
-        *,
-        run_id: UUID,
-        parent_run_id: UUID = None,
-        tags: List[str] = None,
-        **kwargs: Any,
-    ) -> None:
-        settings.chat_history.extend(
-            [
-                AIMessage(content=_parse_token(self.llm,finish.return_values["output"])),
-            ]
-        )
-        finalChunk = {"type": "end"}
-        await self.queue.put(printJson(finalChunk))
-        await self.queue.put(None)
-
-    async def process_json_block(self, json_block: str):
-      """Processa il blocco JSON completo."""
-      # Rimuove il delimitatore iniziale '```json' se presente, e spazi vuoti
-      json_block_clean = json_block.replace('```', '').replace('json', '').strip()
-      # Verifica che il blocco non sia vuoto prima di tentare il parsing
-      if json_block_clean:
-          try:
-              # Prova a fare il parsing del JSON
-              parsed_json = json.loads(json_block_clean)
-              await self.queue.put(printJson(parsed_json))
-          except json.JSONDecodeError as e:
-              # Se il JSON è malformato, logga l'errore
-              raise e
-
-class RawAgentHandler(AsyncCallbackHandler):
-
-    def __init__(self,queue: Queue, llm: str) -> None:
-        super().__init__()
-        self.queue = queue
-        self.llm = llm
-    async def on_llm_start(
-        self,
-        serialized: Dict[str, Any],
-        prompts: List[str],
-        *,
-        run_id: UUID,
-        parent_run_id: UUID = None,
-        tags: List[str] = None,
-        metadata: Dict[str, Any] = None,
-        **kwargs: Any,
-    ) -> None:
-        pass
-
-    """async def on_chat_model_start(self, serialized: Dict[str, Any], messages: List[List[BaseMessage]], *, run_id: UUID = None, parent_run_id = None, tags = None, metadata = None, **kwargs: Any) -> Any:
-        pass"""
-
-    async def on_tool_end(self, output: Any, *, run_id: UUID, parent_run_id: UUID = None, tags: List[str] = None, **kwargs: Any) -> None:
-      pass
-
-    async def on_llm_new_token(
-        self,
-        token: str,
-        *,
-        chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]] = None,
-        run_id: UUID,
-        parent_run_id: Optional[UUID] = None,
-        tags: Optional[List[str]] = None,
-        **kwargs: Any,
-    ) -> None:
-        """Handles new tokens during streaming."""
-        if token:  # Only process non-empty tokens
-            await self.queue.put(_parse_token(self.llm,token))
-
-    async def on_agent_finish(
-        self,
-        finish: AgentFinish,
-        *,
-        run_id: UUID,
-        parent_run_id: UUID = None,
-        tags: List[str] = None,
-        **kwargs: Any,
-    ) -> None:
-        settings.chat_history.extend(
-            [
-                AIMessage(content=_parse_token(self.llm,finish.return_values["output"]))
-            ]
-        )
-        await self.queue.put(None)
+from asyncio import Queue
+from langchain_core.agents import AgentFinish
+from langchain_core.outputs import ChatGenerationChunk, GenerationChunk
+from langchain.callbacks.base import AsyncCallbackHandler
+from ws_bom_robot_app.llm.utils.print import print_json, print_string
+from typing import Any, Dict, List, Optional, Union
+from uuid import UUID
+from langchain_core.callbacks.base import AsyncCallbackHandler
+from langchain_core.outputs import ChatGenerationChunk, GenerationChunk
+from langchain_core.messages import BaseMessage, AIMessage
+import json, logging, re
+
+# Here is a custom handler that will print the tokens to stdout.
+# Instead of printing to stdout you can send the data elsewhere; e.g., to a streaming API response
+
+
+def _parse_token(llm:str,token: str) -> str:
+    """Parses the token based on the LLM provider."""
+    if llm == "anthropic" and isinstance(token, list):
+      first = token[0]
+      if 'text' in first:
+        token = first['text']
+      else:
+        #[{'id': 'toolu_01GGLwJcrQ8PvFMUkQPGu8n7', 'input': {}, 'name': 'document_retriever_xxx', 'type': 'tool_use', 'index': 1}]
+        token = ""
+    return token
+
+class AgentHandler(AsyncCallbackHandler):
+
+    def __init__(self, queue: Queue, llm:str, threadId: str = None) -> None:
+        super().__init__()
+        self._threadId = threadId
+        self.queue = queue
+        self.llm = llm
+        self.__started: bool = False
+        # on new token event
+        self.stream_buffer = ""      # accumulates text that hasn't been processed yet
+        self.in_json_block = False
+        self.json_buffer = ""
+        self.json_start_regex = re.compile(r'(`{1,3}\s*json\b)') # detect a potential json start fence.
+        self.json_end_regex = re.compile(r'(`{1,3})')         # an end fence (one to three backticks).
+        self.stream_cut_last_output_chunk_size = 16  # safe cut last chunk size to output if no markers are found
+    async def on_chat_model_start(self, serialized, messages, *, run_id, parent_run_id = None, tags = None, metadata = None, **kwargs):
+        if not self.__started:
+          self.__started = True
+          firstChunk = {
+              "type": "info",
+              "threadId": self._threadId,
+          }
+          await self.queue.put(print_json(firstChunk))
+
+    async def on_llm_new_token(
+        self,
+        token: str,
+        *,
+        chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]] = None,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        if token and "llm_chain" not in (tags or []):
+            token = _parse_token(self.llm,token)
+            if token:
+              self.stream_buffer += token  # append new data to pending buffer
+              if not self.in_json_block:
+                # search for the start of a json block.
+                start_match = self.json_start_regex.search(self.stream_buffer)
+                if start_match:
+                    start_index = start_match.start()
+                    # everything before the start marker is normal content.
+                    if start_index > 0:
+                        _before = self.stream_buffer[:start_index].replace('`','').strip() # remove eventual preceding backticks.
+                        if _before:
+                            await self.queue.put(print_string(_before))
+                    # remove the start marker from pending.
+                    self.stream_buffer = self.stream_buffer[start_match.end():]
+                    # switch into json mode.
+                    self.in_json_block = True
+                    self.json_buffer = ""
+                else:
+                    # no json start marker found. It might be because the marker is split between chunks.
+                    # to avoid losing potential marker fragments, output what we can safely process:
+                    # if the pending text is long, we output most of it except the last few characters.
+                    if len(self.stream_buffer) > self.stream_cut_last_output_chunk_size:
+                        safe_cut = self.stream_buffer[:-3]
+                        if safe_cut.startswith('`'):
+                          safe_cut = safe_cut[1:]
+                        await self.queue.put(print_string(safe_cut))
+                        self.stream_buffer = self.stream_buffer[-3:]
+              else:
+                  # in json block: look for an end fence.
+                  end_match = self.json_end_regex.search(self.stream_buffer)
+                  if end_match:
+                      end_index = end_match.start()
+                      self.json_buffer += self.stream_buffer[:end_index]
+                      try:
+                          data = json.loads(self.json_buffer.replace('`',''))
+                          await self.queue.put(print_json(data))
+                      except json.JSONDecodeError as e:
+                          logging.error(f"on_token: invalid json: {e} | {self.json_buffer}")
+                      finally:
+                          self.json_buffer = ""
+                      # remove the end fence from pending.
+                      self.stream_buffer = self.stream_buffer[end_match.end():].strip()
+                      self.in_json_block = False
+                  else:
+                      # no end marker found
+                      # accumulate everything and break to wait for more data.
+                      self.json_buffer += self.stream_buffer
+                      self.stream_buffer = ""
+
+    async def on_tool_end(self, output, *, run_id, parent_run_id = None, tags = None, **kwargs):
+        if "stream" in (tags or []):
+            await self.queue.put(print_json(output))
+
+    async def on_agent_finish(
+        self,
+        finish: AgentFinish,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID = None,
+        tags: List[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        # end-of-stream: flush any remaining text
+        if self.in_json_block:
+            try:
+                data = json.loads(self.json_buffer)
+                await self.queue.put(print_json(data))
+            except json.JSONDecodeError as e  :
+                logging.error(f"on_agent_finish: invalid json: {e} | {self.json_buffer}")
+                #await self.queue.put(print_string(self.json_buffer))
+        elif self.stream_buffer:
+            await self.queue.put(print_string(self.stream_buffer))
+
+        finalChunk = {"type": "end"}
+        await self.queue.put(print_json(finalChunk))
+        await self.queue.put(None)
+
+
+class RawAgentHandler(AsyncCallbackHandler):
+
+    def __init__(self,queue: Queue, llm: str) -> None:
+        super().__init__()
+        self.queue = queue
+        self.llm = llm
+
+    async def on_llm_new_token(
+        self,
+        token: str,
+        *,
+        chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]] = None,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handles new tokens during streaming."""
+        if token:  # only process non-empty tokens
+            await self.queue.put(_parse_token(self.llm,token))
+
+    async def on_tool_end(self, output, *, run_id, parent_run_id = None, tags = None, **kwargs):
+        if "stream" in (tags or []):
+            await self.queue.put(print_json(output))
+
+    async def on_agent_finish(
+        self,
+        finish: AgentFinish,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID = None,
+        tags: List[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        await self.queue.put(None)

ws_bom_robot_app/llm/agent_lcel.py

@@ -1,54 +1,107 @@
-from typing import Any
-from langchain.agents import AgentExecutor, create_tool_calling_agent
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.runnables import RunnableLambda
-from langchain_core.tools import render_text_description
-from datetime import datetime
-from ws_bom_robot_app.llm.providers.llm_manager import LlmInterface
-from ws_bom_robot_app.llm.models.api import LlmMessage, LlmRules
-from ws_bom_robot_app.llm.utils.agent_utils import get_rules
-from ws_bom_robot_app.llm.defaut_prompt import default_prompt, tool_prompt
-
-class AgentLcel:
-
-    def __init__(self, llm: LlmInterface, sys_message: str, tools: list, rules: LlmRules = None):
-        self.sys_message = sys_message.format(
-            date_stamp=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-            lang="it",
-        )
-        self.__llm = llm
-        self.__tools = tools
-        self.rules = rules
-        self.embeddings = llm.get_embeddings()
-        self.memory_key = "chat_history"
-        self.__llm_with_tools = llm.get_llm().bind_tools(self.__tools) if len(self.__tools) > 0 else llm.get_llm()
-        self.executor = self.__create_agent()
-
-    async def __create_prompt(self, input: dict) -> ChatPromptTemplate:
-        message : LlmMessage = input["input"]
-        input = message.content
-        rules_prompt = await get_rules(self.embeddings, self.rules, input) if self.rules else ""
-        system = default_prompt + (tool_prompt(render_text_description(self.__tools)) if len(self.__tools)>0 else "") + self.sys_message + rules_prompt
-        return ChatPromptTemplate.from_messages(
-            [
-                (
-                    "system", system
-                ),
-                MessagesPlaceholder(variable_name=self.memory_key),
-                ("user", "{input}"),
-                MessagesPlaceholder(variable_name="agent_scratchpad"),
-            ]
-        )
-
-    def __create_agent(self) -> AgentExecutor:
-      agent: Any = (
-          {
-            "input": lambda x: x["input"],
-            "agent_scratchpad": lambda x: self.__llm.get_formatter(x["intermediate_steps"]),
-            "chat_history": lambda x: x["chat_history"],
-          }
-          | RunnableLambda(self.__create_prompt)
-          | self.__llm_with_tools
-          | self.__llm.get_parser()
-      )
-      return AgentExecutor(agent=agent,tools=self.__tools,verbose=False)
+from typing import Any, Optional
+from langchain.agents import AgentExecutor, create_tool_calling_agent
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.runnables import RunnableLambda
+from langchain_core.tools import render_text_description
+from pydantic import create_model, BaseModel
+import chevron
+from ws_bom_robot_app.llm.agent_context import AgentContext
+from ws_bom_robot_app.llm.providers.llm_manager import LlmInterface
+from ws_bom_robot_app.llm.models.api import LlmMessage, LlmRules
+from ws_bom_robot_app.llm.utils.agent import get_rules
+from ws_bom_robot_app.llm.defaut_prompt import default_prompt, tool_prompt
+
+class AgentLcel:
+
+    def __init__(self, llm: LlmInterface, sys_message: str, sys_context: AgentContext, tools: list, rules: LlmRules = None, json_schema: Optional[dict] = None):
+        self.sys_message = chevron.render(template=sys_message,data=sys_context)
+        self.__llm = llm
+        self.__tools = tools
+        self.rules = rules
+        self.json_schema = json_schema
+        self.embeddings = llm.get_embeddings()
+        self.memory_key: str = "chat_history"
+        self.__llm_with_tools = llm.get_llm().bind_tools(self.__tools) if len(self.__tools) > 0 else llm.get_llm()
+        if self.json_schema:
+            self.__pydantic_schema = self.__create_pydantic_schema()
+        else:
+            self.__pydantic_schema = None
+
+        self.executor = self.__create_agent()
+
+    def __create_pydantic_schema(self) -> type[BaseModel]:
+        """Crea un Pydantic model dinamico dallo schema JSON."""
+        if not self.json_schema:
+            return None
+
+        type_map = {
+            "string": str,
+            "text": str,
+            "number": float,
+            "float": float,
+            "int": int,
+            "integer": int,
+            "bool": bool,
+            "boolean": bool,
+            "list": list,
+            "array": list,
+            "dict": dict,
+            "object": dict,
+        }
+
+        fields: dict[str, tuple[Any, Any]] = {}
+        for k, v in self.json_schema.items():
+            if isinstance(v, str):
+                py_type = type_map.get(v.lower(), str)
+                fields[k] = (py_type, ...)
+            elif isinstance(v, dict):
+                fields[k] = (dict, ...)
+            elif isinstance(v, list):
+                fields[k] = (list, ...)
+            else:
+                fields[k] = (type(v), ...)
+
+        return create_model('JsonSchema', **fields)
+
+    def __get_output_parser(self):
+        return self.__llm.get_parser()
+
+    async def __create_prompt(self, input: dict) -> ChatPromptTemplate:
+        from langchain_core.messages import SystemMessage
+        message : LlmMessage = input[self.memory_key][-1]
+        rules_prompt = await get_rules(self.embeddings, self.rules, message.content) if self.rules else ""
+        system = default_prompt + (tool_prompt(render_text_description(self.__tools)) if len(self.__tools)>0 else "") + self.sys_message + rules_prompt
+
+        # Aggiungi istruzioni per output JSON strutturato se necessario
+        if self.json_schema and self.__pydantic_schema:
+            json_instructions = f"\n\nIMPORTANT: You must format your final response as a JSON object with the following structure:\n"
+            for field_name, field_info in self.__pydantic_schema.model_fields.items():
+                field_type = field_info.annotation.__name__ if hasattr(field_info.annotation, '__name__') else str(field_info.annotation)
+                json_instructions += f"- {field_name}: {field_type}\n"
+            json_instructions += "\nProvide ONLY the JSON object in your response, no additional text."
+            system += json_instructions
+
+        messages = [
+            SystemMessage(content=system),
+            MessagesPlaceholder(variable_name=self.memory_key),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
+        ]
+
+        prompt = ChatPromptTemplate.from_messages(
+            messages=messages,
+            template_format=None,
+        )
+        return prompt
+
+    def __create_agent(self):
+        # Un solo AgentExecutor per entrambe le modalità
+        agent = (
+            {
+                "agent_scratchpad": lambda x: self.__llm.get_formatter(x["intermediate_steps"]),
+                self.memory_key: lambda x: x[self.memory_key],
+            }
+            | RunnableLambda(self.__create_prompt)
+            | self.__llm_with_tools
+            | self.__get_output_parser()
+        )
+        return AgentExecutor(agent=agent, tools=self.__tools, verbose=False)