ws-bom-robot-app 0.0.69__py3-none-any.whl → 0.0.71__py3-none-any.whl

ws_bom_robot_app/llm/nebuly_handler.py

@@ -1,181 +1,181 @@
-from typing import Union
-from ws_bom_robot_app.llm.models.api import NebulyInteraction, NebulyLLMTrace, NebulyRetrievalTrace
-from datetime import datetime, timezone
-from langchain_core.callbacks.base import AsyncCallbackHandler
-import ws_bom_robot_app.llm.settings as settings
-from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
-from langchain_core.outputs import ChatGenerationChunk, GenerationChunk
-
-class NebulyHandler(AsyncCallbackHandler):
-    def __init__(self, llm_model: str | None, threadId: str = None, url: str = None, api_key: str = None, user_id: str | None = None):
-        super().__init__()
-        self.__started: bool = False
-        self.__url: str = url
-        self.__api_key: str = api_key
-        self.interaction = NebulyInteraction(
-            conversation_id=threadId,
-            input="",
-            output="",
-            time_start="",
-            time_end="",
-            end_user= user_id if user_id and user_id != "" else threadId,
-            tags={"model": llm_model},
-        )
-        self.llm_trace = NebulyLLMTrace(
-            model=llm_model,
-            messages=[],
-            output="",
-            input_tokens=0,
-            output_tokens=0,
-        )
-        self.__response_with_rag: str = "false" # Flag to check if the AI used some retrieval tools
-        self.__retrieval_query: str = ""
-        self.retrieval_traces: list[NebulyRetrievalTrace] = []
-
-    async def on_chat_model_start(self, serialized, messages, *, run_id, parent_run_id = None, tags = None, metadata = None, **kwargs):
-        # Initialize the interaction with the input message
-        if not self.__started:
-          message_list = self.__flat_messages(messages)
-          if isinstance(message_list[-1], HumanMessage):
-            if isinstance(message_list[-1].content, list):
-              self.interaction.input = self.__parse_multimodal_input(message_list[-1].content)
-            else:
-              self.interaction.input = message_list[-1].content
-              self.interaction.tags["generated"] = self.__is_message_generated(message_list)
-          else:
-            raise ValueError("Last message is not a HumanMessage")
-          self.interaction.time_start = datetime.now().astimezone().isoformat()
-          self.__started = True
-
-    async def on_llm_end(self, response, *, run_id, parent_run_id = None, tags = None, **kwargs):
-        generation: Union[ChatGenerationChunk, GenerationChunk] = response.generations[0]
-        usage_metadata: dict = generation[0].message.usage_metadata
-        self.llm_trace.input_tokens = usage_metadata.get("input_tokens", 0)
-        self.llm_trace.output_tokens = usage_metadata.get("output_tokens", 0)
-
-    async def on_retriever_start(self, serialized, query, *, run_id, parent_run_id = None, tags = None, metadata = None, **kwargs):
-        self.__retrieval_query = query
-
-
-    async def on_retriever_end(self, documents, *, run_id, parent_run_id = None, tags = None, **kwargs):
-        # pass the document source because of the large amount of data in the document content
-        for doc in documents:
-          self.retrieval_traces.append(
-            NebulyRetrievalTrace(
-              source=doc.metadata.get("source", "content unavailable"),
-              input=self.__retrieval_query,
-              outputs=[doc.metadata.get("source", "content unavailable")]
-            )
-          )
-
-    async def on_tool_start(self, serialized, input_str, *, run_id, parent_run_id = None, tags = None, metadata = None, inputs = None, **kwargs):
-        self.__response_with_rag = "true"  # Set the flag to true when the retriever starts
-
-    async def on_agent_finish(self, finish, *, run_id, parent_run_id = None, tags = None, **kwargs):
-        # Interaction
-        self.interaction.output = finish.return_values["output"]
-        # Trace
-        self.llm_trace.output = finish.return_values["output"]
-        message_history =  self._convert_to_json_format(settings.chat_history)[:-1]
-        self.llm_trace.messages = self.__parse_multimodal_history(message_history)
-        await self.__send_interaction()
-
-    def __flat_messages(self, messages: list[list[BaseMessage]], to_json: bool = False) -> list[BaseMessage]:
-      """
-      Maps the messages to the format expected by the LLM.
-      Flattens the nested list structure of messages.
-      """
-      # Flatten the nested list structure
-      flattened_messages = []
-      for message_list in messages:
-        flattened_messages.extend(message_list)
-      # Store JSON format in LLM trace
-      if to_json:
-        return self._convert_to_json_format(flattened_messages)
-      return flattened_messages
-
-    def _convert_to_json_format(self, messages: list[BaseMessage]) -> list[dict]:
-      """Converts BaseMessage objects to JSON format with role and content."""
-      result = []
-      for message in messages:
-        if isinstance(message, HumanMessage):
-          role = "user"
-        elif isinstance(message, AIMessage):
-          role = "assistant"
-        else:
-          role = "system"
-
-        result.append({
-          "role": role,
-          "content": message.content
-        })
-      return result
-
-    async def __send_interaction(self):
-        # Send the interaction to the server
-        from urllib.parse import urljoin
-        import requests
-
-        payload = self.__prepare_payload()
-        endpoint = urljoin(self.__url, "event-ingestion/api/v2/events/trace_interaction")
-        # Prepare headers with authentication
-        headers = {"Content-Type": "application/json"}
-        if self.__api_key:
-          headers["Authorization"] = f"Bearer {self.__api_key}"
-        response = requests.post(
-          url=endpoint,
-          json=payload,
-          headers=headers
-        )
-        if response.status_code != 200:
-          print(f"Failed to send interaction: {response.status_code} {response.text}")
-
-    def __prepare_payload(self):
-        self.interaction.time_end = datetime.now().astimezone().isoformat()
-        self.interaction.tags["response_with_rag"] = self.__response_with_rag
-        payload = {
-            "interaction": self.interaction.__dict__,
-            "traces": [
-              self.llm_trace.__dict__,
-            ]
-        }
-        for trace in self.retrieval_traces:
-          if trace.source:
-              payload["traces"].append(trace.__dict__)
-        return payload
-
-    def __parse_multimodal_input(self, input: list[dict]) -> str:
-        # Parse the multimodal input and return a string representation
-        # This is a placeholder implementation, you can customize it as needed
-        parsed_input = ""
-        for item in input:
-            if item.get("type") == "text":
-                parsed_input += item.get("text", "")
-            elif item.get("type") == "image_url":
-                parsed_input += " <image>"
-        print(parsed_input)
-        return parsed_input
-
-    def __parse_multimodal_history(self, messages: list[dict]) -> list[dict]:
-        # Parse the multimodal history and return a list of dictionaries
-        parsed_history = []
-        for message in messages:
-            if isinstance(message["content"], list):
-              parsed_content = self.__parse_multimodal_input(message["content"])
-            else:
-              parsed_content = message["content"]
-            parsed_history.append({
-          "role": message["role"],
-          "content": parsed_content
-            })
-        return parsed_history
-
-    def __is_message_generated(self, messages: list[BaseMessage]) -> bool:
-        # Check if the last message is generated by the model
-        if len(messages) == 0:
-            return False
-        last_user_message = f'<div class="llm__pill">{messages[-1].content}</div>'
-        last_ai_message = messages[-2].content
-        if last_user_message in last_ai_message:
-            return "true"
-        return "false"
+from typing import Union
+from ws_bom_robot_app.llm.models.api import NebulyInteraction, NebulyLLMTrace, NebulyRetrievalTrace
+from datetime import datetime, timezone
+from langchain_core.callbacks.base import AsyncCallbackHandler
+import ws_bom_robot_app.llm.settings as settings
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
+from langchain_core.outputs import ChatGenerationChunk, GenerationChunk
+
+class NebulyHandler(AsyncCallbackHandler):
+    def __init__(self, llm_model: str | None, threadId: str = None, url: str = None, api_key: str = None, user_id: str | None = None):
+        super().__init__()
+        self.__started: bool = False
+        self.__url: str = url
+        self.__api_key: str = api_key
+        self.interaction = NebulyInteraction(
+            conversation_id=threadId,
+            input="",
+            output="",
+            time_start="",
+            time_end="",
+            end_user= user_id if user_id and user_id != "" else threadId,
+            tags={"model": llm_model},
+        )
+        self.llm_trace = NebulyLLMTrace(
+            model=llm_model,
+            messages=[],
+            output="",
+            input_tokens=0,
+            output_tokens=0,
+        )
+        self.__response_with_rag: str = "false" # Flag to check if the AI used some retrieval tools
+        self.__retrieval_query: str = ""
+        self.retrieval_traces: list[NebulyRetrievalTrace] = []
+
+    async def on_chat_model_start(self, serialized, messages, *, run_id, parent_run_id = None, tags = None, metadata = None, **kwargs):
+        # Initialize the interaction with the input message
+        if not self.__started:
+          message_list = self.__flat_messages(messages)
+          if isinstance(message_list[-1], HumanMessage):
+            if isinstance(message_list[-1].content, list):
+              self.interaction.input = self.__parse_multimodal_input(message_list[-1].content)
+            else:
+              self.interaction.input = message_list[-1].content
+              self.interaction.tags["generated"] = self.__is_message_generated(message_list)
+          else:
+            raise ValueError("Last message is not a HumanMessage")
+          self.interaction.time_start = datetime.now().astimezone().isoformat()
+          self.__started = True
+
+    async def on_llm_end(self, response, *, run_id, parent_run_id = None, tags = None, **kwargs):
+        generation: Union[ChatGenerationChunk, GenerationChunk] = response.generations[0]
+        usage_metadata: dict = generation[0].message.usage_metadata
+        self.llm_trace.input_tokens = usage_metadata.get("input_tokens", 0)
+        self.llm_trace.output_tokens = usage_metadata.get("output_tokens", 0)
+
+    async def on_retriever_start(self, serialized, query, *, run_id, parent_run_id = None, tags = None, metadata = None, **kwargs):
+        self.__retrieval_query = query
+
+
+    async def on_retriever_end(self, documents, *, run_id, parent_run_id = None, tags = None, **kwargs):
+        # pass the document source because of the large amount of data in the document content
+        for doc in documents:
+          self.retrieval_traces.append(
+            NebulyRetrievalTrace(
+              source=doc.metadata.get("source", "content unavailable"),
+              input=self.__retrieval_query,
+              outputs=[doc.metadata.get("source", "content unavailable")]
+            )
+          )
+
+    async def on_tool_start(self, serialized, input_str, *, run_id, parent_run_id = None, tags = None, metadata = None, inputs = None, **kwargs):
+        self.__response_with_rag = "true"  # Set the flag to true when the retriever starts
+
+    async def on_agent_finish(self, finish, *, run_id, parent_run_id = None, tags = None, **kwargs):
+        # Interaction
+        self.interaction.output = finish.return_values["output"]
+        # Trace
+        self.llm_trace.output = finish.return_values["output"]
+        message_history =  self._convert_to_json_format(settings.chat_history)[:-1]
+        self.llm_trace.messages = self.__parse_multimodal_history(message_history)
+        await self.__send_interaction()
+
+    def __flat_messages(self, messages: list[list[BaseMessage]], to_json: bool = False) -> list[BaseMessage]:
+      """
+      Maps the messages to the format expected by the LLM.
+      Flattens the nested list structure of messages.
+      """
+      # Flatten the nested list structure
+      flattened_messages = []
+      for message_list in messages:
+        flattened_messages.extend(message_list)
+      # Store JSON format in LLM trace
+      if to_json:
+        return self._convert_to_json_format(flattened_messages)
+      return flattened_messages
+
+    def _convert_to_json_format(self, messages: list[BaseMessage]) -> list[dict]:
+      """Converts BaseMessage objects to JSON format with role and content."""
+      result = []
+      for message in messages:
+        if isinstance(message, HumanMessage):
+          role = "user"
+        elif isinstance(message, AIMessage):
+          role = "assistant"
+        else:
+          role = "system"
+
+        result.append({
+          "role": role,
+          "content": message.content
+        })
+      return result
+
+    async def __send_interaction(self):
+        # Send the interaction to the server
+        from urllib.parse import urljoin
+        import requests
+
+        payload = self.__prepare_payload()
+        endpoint = urljoin(self.__url, "event-ingestion/api/v2/events/trace_interaction")
+        # Prepare headers with authentication
+        headers = {"Content-Type": "application/json"}
+        if self.__api_key:
+          headers["Authorization"] = f"Bearer {self.__api_key}"
+        response = requests.post(
+          url=endpoint,
+          json=payload,
+          headers=headers
+        )
+        if response.status_code != 200:
+          print(f"Failed to send interaction: {response.status_code} {response.text}")
+
+    def __prepare_payload(self):
+        self.interaction.time_end = datetime.now().astimezone().isoformat()
+        self.interaction.tags["response_with_rag"] = self.__response_with_rag
+        payload = {
+            "interaction": self.interaction.__dict__,
+            "traces": [
+              self.llm_trace.__dict__,
+            ]
+        }
+        for trace in self.retrieval_traces:
+          if trace.source:
+              payload["traces"].append(trace.__dict__)
+        return payload
+
+    def __parse_multimodal_input(self, input: list[dict]) -> str:
+        # Parse the multimodal input and return a string representation
+        # This is a placeholder implementation, you can customize it as needed
+        parsed_input = ""
+        for item in input:
+            if item.get("type") == "text":
+                parsed_input += item.get("text", "")
+            elif item.get("type") == "image_url":
+                parsed_input += " <image>"
+        print(parsed_input)
+        return parsed_input
+
+    def __parse_multimodal_history(self, messages: list[dict]) -> list[dict]:
+        # Parse the multimodal history and return a list of dictionaries
+        parsed_history = []
+        for message in messages:
+            if isinstance(message["content"], list):
+              parsed_content = self.__parse_multimodal_input(message["content"])
+            else:
+              parsed_content = message["content"]
+            parsed_history.append({
+          "role": message["role"],
+          "content": parsed_content
+            })
+        return parsed_history
+
+    def __is_message_generated(self, messages: list[BaseMessage]) -> bool:
+        # Check if the last message is generated by the model
+        if len(messages) == 0:
+            return False
+        last_user_message = f'<div class="llm__pill">{messages[-1].content}</div>'
+        last_ai_message = messages[-2].content
+        if last_user_message in last_ai_message:
+            return "true"
+        return "false"

ws_bom_robot_app/llm/providers/llm_manager.py

@@ -107,9 +107,8 @@ class Google(LlmInterface):
   def get_llm(self):
     from langchain_google_genai.chat_models import ChatGoogleGenerativeAI
     return ChatGoogleGenerativeAI(
-      name="chat",
-      api_key=self.config.api_key or os.getenv("GOOGLE_API_KEY"),
       model=self.config.model,
+      google_api_key=self.config.api_key or os.getenv("GOOGLE_API_KEY"),
       temperature=self.config.temperature,
       disable_streaming=False,
     )
@@ -117,8 +116,8 @@ class Google(LlmInterface):
   def get_embeddings(self):
     from langchain_google_genai.embeddings import GoogleGenerativeAIEmbeddings
     return GoogleGenerativeAIEmbeddings(
-      google_api_key=self.config.api_key,
-      model="models/text-embedding-005")
+      google_api_key=self.config.api_key or os.getenv("GOOGLE_API_KEY"),
+      model="models/text-embedding-004")
 
   def get_models(self):
     import google.generativeai as genai
@@ -140,22 +139,33 @@ class Gvertex(LlmInterface):
             temperature=self.config.temperature
         )
     def get_embeddings(self):
-        from langchain_google_vertexai import VertexAIEmbeddings
-        return VertexAIEmbeddings(model_name="text-embedding-005")
+        from langchain_google_vertexai.embeddings import VertexAIEmbeddings
+        embeddings = VertexAIEmbeddings(model_name="text-embedding-004")
+        #fix gemini-embedding-001 batch size
+        #embeddings.instance["max_batch_size"] = 1
+        #embeddings.instance["batch_size"] = 1
+        return embeddings
     def get_models(self):
-        #from google.cloud import aiplatform
-        #aiplatform.init()
-        #models = aiplatform.Model.list()
+        _models = [
+                {"id":"gemini-2.5-pro"},
+                {"id":"gemini-2.5-flash"},
+                {"id":"gemini-2.0-flash"},
+                {"id":"gemini-2.0-flash-lite"}
+              ]
+        try:
+          from google.cloud import aiplatform
+          aiplatform.init()
+          _list = aiplatform.Model.list()
+          if _list:
+            _models = list([{"id": model.name} for model in _list])
         # removed due issue: https://github.com/langchain-ai/langchain-google/issues/733
         # Message type "google.cloud.aiplatform.v1beta1.GenerateContentResponse" has no field named "createTime" at "GenerateContentResponse".  Available Fields(except extensions): "['candidates', 'modelVersion', 'promptFeedback', 'usageMetadata']"
-
-        #see https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#united-states for available models
-        return [
-              {"id":"gemini-2.5-pro-preview-05-06"},
-              {"id":"gemini-2.0-flash"},
-              {"id":"gemini-2.0-flash-lite"},
-              {"id":"gemini-1.5-pro-002"}
-            ]
+        except Exception as e:
+          print(f"Error fetching models from Gvertex: {e}")
+          # fallback to hardcoded models
+          #see https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#united-states for available models
+        finally:
+          return _models
 
 class Groq(LlmInterface):
     def get_llm(self):

ws_bom_robot_app/llm/settings.py

@@ -1,4 +1,4 @@
-def init():
-    """Initialize the chat history list as a global var"""
-    global chat_history
-    chat_history = []
+def init():
+    """Initialize the chat history list as a global var"""
+    global chat_history
+    chat_history = []

ws_bom_robot_app/llm/tools/tool_builder.py

@@ -1,65 +1,65 @@
-import asyncio
-from asyncio import Queue
-from langchain.tools import StructuredTool
-from ws_bom_robot_app.llm.models.api import LlmAppTool
-from ws_bom_robot_app.llm.tools.tool_manager import ToolManager
-from ws_bom_robot_app.llm.providers.llm_manager import LlmInterface
-
-async def __process_proxy_tool(proxy_tool: LlmAppTool) -> LlmAppTool | None:
-    import os
-    from ws_bom_robot_app.llm.utils.cms import CmsApp, get_app_by_id
-    from ws_bom_robot_app.config import config
-    try:
-        secrets = proxy_tool.secrets_to_dict()
-        app_id = secrets.get("appId")
-        if not app_id:
-            raise ValueError("Tool configuration is invalid. 'appId' is required.")
-        app: CmsApp = await get_app_by_id(app_id)
-        if not app:
-            raise ValueError(f"App with id {app_id} not found.")
-        tool_id = secrets.get("toolId")
-        tool = next((t for t in app.rq.app_tools if app.rq.app_tools and t.id == tool_id), None)
-        if not tool:
-            raise ValueError(f"Tool with function_id {tool_id} not found in app {app.name}.")
-        #override derived tool with proxy tool props
-        tool.name = proxy_tool.name if proxy_tool.name else tool.name
-        tool.description = proxy_tool.description if proxy_tool.description else tool.description
-        tool.function_id = proxy_tool.function_id if proxy_tool.function_id else tool.function_id
-        tool.function_description = proxy_tool.function_description if proxy_tool.function_description else tool.function_description
-        #normalize vector_db
-        if tool.vector_db:
-            tool.vector_db = os.path.join(
-               os.path.join(config.robot_data_folder,config.robot_data_db_folder,config.robot_data_db_folder_store),
-               os.path.splitext(os.path.basename(tool.vector_db))[0]) if tool.vector_db else None
-        return tool
-    except Exception as e:
-        print(f"[!] Error in proxy_app_tool: {e}")
-        return None
-
-def get_structured_tools(llm: LlmInterface, tools: list[LlmAppTool], callbacks:list, queue: Queue) -> list[StructuredTool]:
-  _structured_tools :list[StructuredTool] = []
-  for tool in [tool for tool in tools if tool.is_active]:
-    if tool.function_name == "proxy_app_tool":
-        # override the tool
-        loop = asyncio.get_event_loop()
-        if loop.is_running():
-          import nest_asyncio
-          nest_asyncio.apply()
-        processed_tool = loop.run_until_complete(__process_proxy_tool(tool))
-        if processed_tool is None:
-            continue
-        tool = processed_tool
-    if _tool_config := ToolManager._list.get(tool.function_name):
-      _tool_instance = ToolManager(llm, tool, callbacks, queue)
-      _structured_tool = StructuredTool.from_function(
-        coroutine=_tool_instance.get_coroutine(),
-        name=tool.function_id if tool.function_id else tool.function_name,
-        description=tool.function_description,
-        args_schema=_tool_config.model
-        #infer_schema=True,
-        #parse_docstring=True,
-        #error_on_invalid_docstring=True
-      )
-      _structured_tool.tags = [tool.function_id if tool.function_id else tool.function_name]
-      _structured_tools.append(_structured_tool)
-  return _structured_tools
+import asyncio
+from asyncio import Queue
+from langchain.tools import StructuredTool
+from ws_bom_robot_app.llm.models.api import LlmAppTool
+from ws_bom_robot_app.llm.tools.tool_manager import ToolManager
+from ws_bom_robot_app.llm.providers.llm_manager import LlmInterface
+
+async def __process_proxy_tool(proxy_tool: LlmAppTool) -> LlmAppTool | None:
+    import os
+    from ws_bom_robot_app.llm.utils.cms import CmsApp, get_app_by_id
+    from ws_bom_robot_app.config import config
+    try:
+        secrets = proxy_tool.secrets_to_dict()
+        app_id = secrets.get("appId")
+        if not app_id:
+            raise ValueError("Tool configuration is invalid. 'appId' is required.")
+        app: CmsApp = await get_app_by_id(app_id)
+        if not app:
+            raise ValueError(f"App with id {app_id} not found.")
+        tool_id = secrets.get("toolId")
+        tool = next((t for t in app.rq.app_tools if app.rq.app_tools and t.id == tool_id), None)
+        if not tool:
+            raise ValueError(f"Tool with function_id {tool_id} not found in app {app.name}.")
+        #override derived tool with proxy tool props
+        tool.name = proxy_tool.name if proxy_tool.name else tool.name
+        tool.description = proxy_tool.description if proxy_tool.description else tool.description
+        tool.function_id = proxy_tool.function_id if proxy_tool.function_id else tool.function_id
+        tool.function_description = proxy_tool.function_description if proxy_tool.function_description else tool.function_description
+        #normalize vector_db
+        if tool.vector_db:
+            tool.vector_db = os.path.join(
+               os.path.join(config.robot_data_folder,config.robot_data_db_folder,config.robot_data_db_folder_store),
+               os.path.splitext(os.path.basename(tool.vector_db))[0]) if tool.vector_db else None
+        return tool
+    except Exception as e:
+        print(f"[!] Error in proxy_app_tool: {e}")
+        return None
+
+def get_structured_tools(llm: LlmInterface, tools: list[LlmAppTool], callbacks:list, queue: Queue) -> list[StructuredTool]:
+  _structured_tools :list[StructuredTool] = []
+  for tool in [tool for tool in tools if tool.is_active]:
+    if tool.function_name == "proxy_app_tool":
+        # override the tool
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+          import nest_asyncio
+          nest_asyncio.apply()
+        processed_tool = loop.run_until_complete(__process_proxy_tool(tool))
+        if processed_tool is None:
+            continue
+        tool = processed_tool
+    if _tool_config := ToolManager._list.get(tool.function_name):
+      _tool_instance = ToolManager(llm, tool, callbacks, queue)
+      _structured_tool = StructuredTool.from_function(
+        coroutine=_tool_instance.get_coroutine(),
+        name=tool.function_id if tool.function_id else tool.function_name,
+        description=tool.function_description,
+        args_schema=_tool_config.model
+        #infer_schema=True,
+        #parse_docstring=True,
+        #error_on_invalid_docstring=True
+      )
+      _structured_tool.tags = [tool.function_id if tool.function_id else tool.function_name]
+      _structured_tools.append(_structured_tool)
+  return _structured_tools