langchain-timbr 1.5.0__py3-none-any.whl

langchain_timbr/llm_wrapper/llm_wrapper.py

@@ -0,0 +1,189 @@
+from enum import Enum
+from langchain.llms.base import LLM
+from pydantic import Field
+
+from .timbr_llm_wrapper import TimbrLlmWrapper
+from ..utils.general import is_llm_type, is_support_temperature
+from ..config import llm_temperature
+
+class LlmTypes(Enum):
+  OpenAI = 'openai-chat'
+  Anthropic = 'anthropic-chat'
+  Google = 'chat-google-generative-ai'
+  AzureOpenAI = 'azure-openai-chat'
+  Snowflake = 'snowflake-cortex'
+  Timbr = 'timbr'
+
+
+class LlmWrapper(LLM):
+  """
+  LlmWrapper is a unified interface for connecting to various Large Language Model (LLM) providers
+  (OpenAI, Anthropic, Google, Azure OpenAI, Snowflake Cortex, etc.) using LangChain. It abstracts
+  the initialization and connection logic for each provider, allowing you to switch between them
+  with a consistent API.
+  """
+  client: LLM = Field(default=None, exclude=True)
+
+  def __init__(
+      self,
+      llm_type: str,
+      api_key: str,
+      model: str = None,
+      **llm_params,
+  ):
+      """
+      :param llm_type (str): The type of LLM provider (e.g., 'openai-chat', 'anthropic-chat').
+      :param api_key (str): The API key for authenticating with the LLM provider.
+      :param model (str): The model name or deployment to use. Defaults to provider-specific values (Optional).
+      :param **llm_params: Additional parameters for the LLM (e.g., temperature, endpoint, etc.).
+      """
+      super().__init__()
+      self.client = self._connect_to_llm(
+        llm_type,
+        api_key,
+        model,
+        **llm_params,
+      )
+
+
+  @property
+  def _llm_type(self):
+    return self.client._llm_type
+
+
+  def _add_temperature(self, llm_type, llm_model, **llm_params):
+    """
+    Add temperature to the LLM parameters if the LLM model supports it.
+    """
+    if "temperature" not in llm_params:
+      if llm_temperature is not None and is_support_temperature(llm_type, llm_model):
+        llm_params["temperature"] = llm_temperature
+    return llm_params
+
+  
+  def _connect_to_llm(self, llm_type, api_key, model, **llm_params):
+    if is_llm_type(llm_type, LlmTypes.OpenAI):
+      from langchain_openai import ChatOpenAI as OpenAI
+      llm_model = model or "gpt-4o-2024-11-20"
+      params = self._add_temperature(LlmTypes.OpenAI.name, llm_model, **llm_params)
+      return OpenAI(
+        openai_api_key=api_key,
+        model_name=llm_model,
+        **params,
+      )
+    elif is_llm_type(llm_type, LlmTypes.Anthropic):
+      from langchain_anthropic import ChatAnthropic as Claude
+      llm_model = model or "claude-3-5-sonnet-20241022"
+      params = self._add_temperature(LlmTypes.Anthropic.name, llm_model, **llm_params)
+      return Claude(
+        anthropic_api_key=api_key,
+        model=llm_model,
+        **params,
+      )
+    elif is_llm_type(llm_type, LlmTypes.Google):
+      from langchain_google_genai import ChatGoogleGenerativeAI
+      llm_model = model or "gemini-2.0-flash-exp"
+      params = self._add_temperature(LlmTypes.Google.name, llm_model, **llm_params)
+      return ChatGoogleGenerativeAI(
+        google_api_key=api_key,
+        model=llm_model,
+        **params,
+      )
+    elif is_llm_type(llm_type, LlmTypes.Timbr):
+      return TimbrLlmWrapper(
+        api_key=api_key,
+        **params,
+      )
+    elif is_llm_type(llm_type, LlmTypes.Snowflake):
+      from langchain_community.chat_models import ChatSnowflakeCortex      
+      llm_model = model or "openai-gpt-4.1"
+      params = self._add_temperature(LlmTypes.Snowflake.name, llm_model, **llm_params)
+      
+      return ChatSnowflakeCortex(
+        model=llm_model,
+        **params,
+      )
+    elif is_llm_type(llm_type, LlmTypes.AzureOpenAI):
+      from langchain_openai import AzureChatOpenAI
+      azure_endpoint = params.pop('azure_endpoint', None)
+      azure_api_version = params.pop('azure_openai_api_version', None)
+      llm_model = model or "gpt-4o-2024-11-20"
+      params = self._add_temperature(LlmTypes.AzureOpenAI.name, llm_model, **llm_params)
+      return AzureChatOpenAI(
+        openai_api_key=api_key,
+        azure_deployment=llm_model,
+        azure_endpoint=azure_endpoint,
+        openai_api_version=azure_api_version,
+        **params,
+      )
+    else:
+      raise ValueError(f"Unsupported LLM type: {llm_type}")
+
+
+  def get_model_list(self) -> list[str]:
+    """Return the list of available models for the LLM."""
+    models = []
+    try:
+      if is_llm_type(self._llm_type, LlmTypes.OpenAI):
+        from openai import OpenAI
+        client = OpenAI(api_key=self.client.openai_api_key._secret_value)
+        models = [model.id for model in client.models.list()]
+      elif is_llm_type(self._llm_type, LlmTypes.Anthropic):
+        import anthropic
+        client = anthropic.Anthropic(api_key=self.client.anthropic_api_key._secret_value)
+        models = [model.id for model in client.models.list()]
+      elif is_llm_type(self._llm_type, LlmTypes.Google):
+        import google.generativeai as genai
+        genai.configure(api_key=self.client.google_api_key._secret_value)
+        models = [m.name.replace('models/', '') for m in genai.list_models()]
+      elif is_llm_type(self._llm_type, LlmTypes.AzureOpenAI):
+        from openai import AzureOpenAI
+        # Get Azure-specific attributes from the client
+        azure_endpoint = getattr(self.client, 'azure_endpoint', None)
+        api_version = getattr(self.client, 'openai_api_version', None)
+        api_key = self.client.openai_api_key._secret_value
+        
+        if azure_endpoint and api_version and api_key:
+          client = AzureOpenAI(
+            api_key=api_key,
+            azure_endpoint=azure_endpoint,
+            api_version=api_version
+          )
+          # For Azure, get the deployments instead of models
+          try:
+            models = [model.id for model in client.models.list()]
+          except:
+            # If listing models fails, provide some common deployment names
+            models = ["gpt-4o", "Other (Custom)"]
+      elif is_llm_type(self._llm_type, LlmTypes.Snowflake):
+        # Snowflake Cortex available models
+        models = [
+          "openai-gpt-4.1",
+          "mistral-large2",
+          "llama3.1-70b",
+          "llama3.1-405b"
+        ]
+      # elif self._is_llm_type(self._llm_type, LlmTypes.Timbr):
+        
+    except Exception as e:
+      models = []
+    
+    return models
+
+
+  def _call(self, prompt, **kwargs):
+    return self.client(prompt, **kwargs)
+
+
+  def __call__(self, prompt, **kwargs):
+        """
+        Override the default __call__ method to handle input preprocessing.
+        I used this in order to override prompt input validation made by pydantic
+        and allow sending list of AiMessages instead of string only
+        """
+        return self._call(prompt, **kwargs)
+  
+
+  def query(self, prompt, **kwargs):
+    return self._call(prompt, **kwargs)
+

langchain_timbr/llm_wrapper/timbr_llm_wrapper.py

@@ -0,0 +1,41 @@
+from langchain.llms.base import LLM
+import requests
+from typing import Optional, List
+
+class TimbrLlmWrapper(LLM):
+    def __init__(self, url: str, api_key: str, temperature: Optional[float] = 0):
+        """
+        ***TBD, Not ready yet.***
+        
+        Custom LLM implementation for timbr LLM wrapped with a proxy server.
+        
+        :param url: URL of the proxy server wrapping timbr LLM.
+        :param api_key: API key for authentication with the proxy server.
+        :param temperature: Sampling temperature for the model.
+        """
+        self.url = url
+        self.api_key = api_key
+        self.temperature = temperature
+
+    @property
+    def _llm_type(self) -> str:
+        return "timbr"
+
+
+    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        """
+        Sends the prompt to the proxy server and returns the response.
+        """
+        headers = { "Authorization": f"Bearer {self.api_key}" }
+        payload = {
+            "prompt": prompt,
+            "temperature": self.temperature,
+            "stop": stop,
+        }
+        
+        response = requests.post(self.url, json=payload, headers=headers)
+        if response.status_code == 200:
+            return response.json().get("response", "")
+        else:
+            raise ValueError(f"Error communicating with timbr proxy: {response.text}")
+

langchain_timbr/timbr_llm_connector.py

@@ -0,0 +1,398 @@
+from typing import Optional, Any, Literal
+from typing_extensions import TypedDict
+from langchain.llms.base import LLM
+from langgraph.graph import StateGraph, END
+
+from .utils.general import to_boolean, to_integer
+from .llm_wrapper.llm_wrapper import LlmWrapper
+from .utils.timbr_utils import get_ontologies, get_concepts
+from .langchain import IdentifyTimbrConceptChain, GenerateTimbrSqlChain, ValidateTimbrSqlChain, ExecuteTimbrQueryChain, create_timbr_sql_agent
+from .langgraph import GenerateTimbrSqlNode, ValidateSemanticSqlNode, ExecuteSemanticQueryNode, GenerateResponseNode
+
+
+from .config import (
+    url as default_url,
+    token as default_token,
+    ontology as default_ontology,
+    llm_type,
+    llm_model,
+    llm_api_key,
+    llm_temperature,
+)
+
+class TimbrLanggraphState(TypedDict):
+    prompt: str
+    sql: str
+    concept: str
+    rows: list
+    response: str
+    error: str
+    is_sql_valid: bool
+    usage_metadata: dict[str, Any]
+
+
+class TimbrLlmConnector:
+    def __init__(
+        self,
+        llm: LLM,
+        url: Optional[str] = default_url,
+        token: Optional[str] = default_token,
+        ontology: Optional[str] = default_ontology,
+        max_limit: Optional[int] = 500,
+        verify_ssl: Optional[bool] = True,
+        is_jwt: Optional[bool] = False,
+        jwt_tenant_id: Optional[str] = None,
+        conn_params: Optional[dict] = None,
+    ):
+        """
+        :param url: Timbr server url
+        :param token: Timbr password or token value
+        :param ontology: The name of the ontology/knowledge graph
+        :param llm: An LLM instance or a function that takes a prompt string and returns the LLM’s response
+        :param max_limit: Maximum number of rows to return
+        :param verify_ssl: Whether to verify SSL certificates (default is True).
+        :param is_jwt: Whether to use JWT authentication (default is False).
+        :param jwt_tenant_id: Tenant ID for JWT authentication (if applicable).
+        :param conn_params: Extra Timbr connection parameters sent with every request (e.g., 'x-api-impersonate-user').
+
+        ## Example
+        ```
+        timbr_llm_wrapper = LlmWrapper(
+            llm_type=LlmTypes.OpenAI,
+            model="gpt-4o"
+            api_key=<openai_api_key>
+        )
+
+        llm_connector = TimbrLlmConnector(
+            url=<url>,
+            token=<token>,
+            llm=timbr_llm_wrapper,
+        )
+
+        # Show ontology list at timbr instance from url connection
+        ontologies = llm_connector.get_ontologies()
+
+        # Find which concept & schema will be queried by the user input
+        determine_concept_res = llm_connector.determine_concept(llm_input)
+        query_concept, query_schema = determine_concept_res.get("concept"), determine_concept_res.get("schema")
+
+        # Generate timbr SQL query from user input
+        sql_query = llm_connector.generate_sql(llm_input).get("sql")
+
+        # Run timbr SQL query
+        results = llm_connector.run_timbr_query(sql_query).get("rows", [])
+        
+        # Parse & Run LLM question
+        results = llm_connector.run_llm_query(llm_input).get("rows", [])
+        ```
+        """
+        self.url = url
+        self.token = token
+        self.ontology = ontology
+        self.max_limit = to_integer(max_limit)
+        self.verify_ssl = to_boolean(verify_ssl)
+        self.is_jwt = to_boolean(is_jwt)
+        self.jwt_tenant_id = jwt_tenant_id
+        self.conn_params = conn_params or {}
+        
+        if llm is not None:
+            self._llm = llm
+        elif llm_type is not None and llm_api_key is not None:
+            llm_params = {}
+            if llm_temperature is not None:
+                llm_params["temperature"] = llm_temperature
+
+            self._llm = LlmWrapper(
+                llm_type=llm_type,
+                api_key=llm_api_key,
+                model=llm_model,
+                **llm_params,
+            )
+
+
+    # TODO: Make this function a decorator and use in on relevant methods
+    # def _is_ontology_set(self):
+    #     return self.ontology != 'system_db'
+    
+
+    def _get_conn_params(self):
+        return {
+            "url": self.url,
+            "token": self.token,
+            "ontology": self.ontology,
+            "verify_ssl": self.verify_ssl,
+            "is_jwt": self.is_jwt,
+            "jwt_tenant_id": self.jwt_tenant_id,
+            **self.conn_params,
+        }
+
+
+    def get_ontologies(self) -> list[str]:
+        return get_ontologies(conn_params=self._get_conn_params())
+
+
+    def get_concepts(self) -> dict:
+        """
+        Get the list of concepts from the Timbr server.
+        """
+        return get_concepts(
+            conn_params=self._get_conn_params(),
+            concepts_list="*",
+        )
+    
+
+    def get_views(self) -> dict:
+        """
+        Get the list of views from the Timbr server.
+        """
+        return get_concepts(
+            conn_params=self._get_conn_params(),
+            views_list="*",
+        )
+
+
+    def set_ontology(self, ontology: str):
+        self.ontology = ontology
+
+
+    def determine_concept(
+        self,
+        question: str,
+        concepts_list: Optional[list] = None,
+        views_list: Optional[list] = None,
+        include_logic_concepts: Optional[bool] = False,
+        include_tags: Optional[str] = None,
+        should_validate: Optional[bool] = False,
+        retries: Optional[int] = 3,
+        note: Optional[str] = '',
+        **chain_kwargs: Any,
+    ) -> dict[str, Any]:
+        determine_concept_chain = IdentifyTimbrConceptChain(
+            **self._get_conn_params(),
+            llm=self._llm,
+            concepts_list=concepts_list,
+            views_list=views_list,
+            include_logic_concepts=include_logic_concepts,
+            include_tags=include_tags,
+            should_validate=should_validate,
+            retries=retries,
+            note=note,
+            **chain_kwargs,
+        )
+
+        return determine_concept_chain.invoke({ "prompt": question })
+
+
+    def generate_sql(
+        self,
+        question: str,
+        concept_name: Optional[str] = None,
+        schema: Optional[str] = None,
+        concepts_list: Optional[list] = None,
+        views_list: Optional[list] = None,
+        include_logic_concepts: Optional[bool] = False,
+        include_tags: Optional[str] = None,
+        should_validate_sql: Optional[bool] = False,
+        retries: Optional[int] = 3,
+        note: Optional[str] = '',
+        **chain_kwargs: Any,
+    ) -> dict[str, Any]:
+        generate_timbr_llm_chain = GenerateTimbrSqlChain(
+            llm=self._llm,
+            **self._get_conn_params(),
+            schema=schema,
+            concept=concept_name,
+            concepts_list=concepts_list,
+            views_list=views_list,
+            include_logic_concepts=include_logic_concepts,
+            include_tags=include_tags,
+            should_validate_sql=should_validate_sql,
+            retries=retries,
+            max_limit=self.max_limit,
+            note=note,
+            **chain_kwargs,
+        )
+
+        return generate_timbr_llm_chain.invoke({ "prompt": question })
+
+
+    def validate_sql(
+        self,
+        question: str,
+        sql_query: str,
+        retries: Optional[int] = 3,
+        concepts_list: Optional[list] = None,
+        views_list: Optional[list] = None,
+        include_logic_concepts: Optional[bool] = False,
+        include_tags: Optional[str] = None,
+        note: Optional[str] = '',
+        **chain_kwargs: Any,
+    ) -> dict[str, Any]:
+        validate_timbr_sql_chain = ValidateTimbrSqlChain(
+            llm=self._llm,
+            **self._get_conn_params(),
+            retries=retries,
+            concepts_list=concepts_list,
+            views_list=views_list,
+            include_logic_concepts=include_logic_concepts,
+            include_tags=include_tags,
+            max_limit=self.max_limit,
+            note=note,
+            **chain_kwargs,
+        )
+        return validate_timbr_sql_chain.invoke({ "sql": sql_query, "prompt": question })
+
+
+    def run_timbr_query(
+        self,
+        sql_query: str,
+        concepts_list: Optional[list] = None,
+        views_list: Optional[list] = None,
+        include_logic_concepts: Optional[bool] = False,
+        include_tags: Optional[str] = None,
+        should_validate_sql: Optional[bool] = True,
+        retries: Optional[int] = 3,
+        note: Optional[str] = '',
+        **chain_kwargs: Any,
+    ) -> dict[str, Any]:
+        execute_timbr_query_chain = ExecuteTimbrQueryChain(
+            llm=self._llm,
+            **self._get_conn_params(),
+            concepts_list=concepts_list,
+            views_list=views_list,
+            include_logic_concepts=include_logic_concepts,
+            include_tags=include_tags,
+            should_validate_sql=should_validate_sql,
+            retries=retries,
+            max_limit=self.max_limit,
+            note=note,
+            **chain_kwargs,
+        )
+
+        return execute_timbr_query_chain.invoke({ "sql": sql_query })
+
+
+    def run_llm_query(
+        self,
+        question: str,
+        concepts_list: Optional[list] = None,
+        views_list: Optional[list] = None,
+        include_logic_concepts: Optional[bool] = False,
+        include_tags: Optional[str] = None,
+        should_validate_sql: Optional[bool] = True,
+        retries: Optional[int] = 3,
+        note: Optional[str] = '',
+        **agent_kwargs: Any,
+    ) -> dict[str, Any]:
+        agent = create_timbr_sql_agent(
+            llm=self._llm,
+            **self._get_conn_params(),
+            concept=None,
+            concepts_list=concepts_list,
+            views_list=views_list,
+            include_logic_concepts=include_logic_concepts,
+            include_tags=include_tags,
+            should_validate_sql=should_validate_sql,
+            retries=retries,
+            max_limit=self.max_limit,
+            note=note,
+            **agent_kwargs,
+        )
+
+        return agent.invoke(question)
+
+
+    def run_llm_query_graph(
+        self,
+        question: str,
+        concepts_list: Optional[list] = None,
+        views_list: Optional[list] = None,
+        include_logic_concepts: Optional[bool] = False,
+        include_tags: Optional[str] = None,
+        should_validate_sql: Optional[bool] = True,
+        retries: Optional[int] = 3,
+        note: Optional[str] = '',
+        **nodes_kwargs: Any,
+    ) -> dict[str, Any]:
+        generate_sql_node = GenerateTimbrSqlNode(
+            llm=self._llm,
+            **self._get_conn_params(),
+            concepts_list=concepts_list,
+            views_list=views_list,
+            include_logic_concepts=include_logic_concepts,
+            include_tags=include_tags,
+            max_limit=self.max_limit,
+            note=note,
+            **nodes_kwargs,
+        )
+        validate_sql_node = ValidateSemanticSqlNode(
+            llm=self._llm,
+            **self._get_conn_params(),
+            retries=retries,
+            concepts_list=concepts_list,
+            views_list=views_list,
+            include_logic_concepts=include_logic_concepts,
+            include_tags=include_tags,
+            max_limit=self.max_limit,
+            note=note,
+            **nodes_kwargs,
+        )
+        execute_sql_node = ExecuteSemanticQueryNode(
+            llm=self._llm,
+            **self._get_conn_params(),
+            concepts_list=concepts_list,
+            views_list=views_list,
+            include_logic_concepts=include_logic_concepts,
+            include_tags=include_tags,
+            should_validate_sql=should_validate_sql,
+            retries=retries,
+            max_limit=self.max_limit,
+            note=note,
+            **nodes_kwargs,
+        )
+        generate_response_node = GenerateResponseNode()
+
+        graph_builder = StateGraph(TimbrLanggraphState)
+
+        graph_builder.add_node("generate_sql", generate_sql_node)
+        graph_builder.add_node("validate_sql", validate_sql_node)
+        graph_builder.add_node("execute_sql", execute_sql_node)
+        graph_builder.add_node("generate_response", generate_response_node)
+
+        graph_builder.add_edge("generate_sql", "validate_sql")
+        
+        def route_validation(state: dict) -> Literal["execute_sql", "end"]:
+            # If validation is successful, proceed to execute the query.
+            # Otherwise, stop the flow.
+            if state.get("is_sql_valid"):
+                return "execute_sql"
+            else:
+                return "end"
+            
+        graph_builder.add_conditional_edges(
+            "validate_sql", 
+            route_validation,
+            {
+                "execute_sql": "execute_sql",
+                "end": END
+            }
+        )
+        
+        graph_builder.add_edge("execute_sql", "generate_response")
+        graph_builder.set_entry_point("generate_sql")
+
+        compiled_graph = graph_builder.compile()
+        
+        initial_state = {
+            "prompt": question,
+            "sql": "",
+            "concept": "",
+            "rows": [],
+            "response": "",
+            "error": "",
+            "is_sql_valid": False,
+            "usage_metadata": {}
+        }
+        
+        result = compiled_graph.invoke(initial_state)
+        return result

langchain_timbr/utils/general.py

@@ -0,0 +1,70 @@
+import os
+import json
+
+### A global helper functions to use across the project
+
+def parse_list(input_value, separator=',') -> list[str]:
+    try:
+        if isinstance(input_value, str):
+            return [item.strip() for item in input_value.split(separator) if item.strip()]
+        elif isinstance(input_value, list):
+            return [item.strip() for item in input_value if item.strip()]
+        return []
+    except Exception as e:
+        raise ValueError(f"Failed to parse list value: {e}")
+
+
+def to_boolean(value) -> bool:
+    try:
+        if isinstance(value, str):
+            return value.lower() in ['true', '1']
+        return bool(value)
+    except Exception as e:
+        raise ValueError(f"Failed to parse boolean value: {e}")
+
+
+def to_integer(value) -> int:
+    try:
+        return int(value)
+    except (ValueError, TypeError) as e:
+        raise ValueError(f"Failed to parse integer value: {e}")
+
+
+def is_llm_type(llm_type, enum_value):
+    """Check if llm_type equals the enum value or its name, case-insensitive."""
+    if llm_type == enum_value:
+        return True
+    
+    if isinstance(llm_type, str):
+        llm_type_lower = llm_type.lower()
+        enum_name_lower = enum_value.name.lower() if enum_value.name else ""
+        enum_value_lower = enum_value.value.lower() if isinstance(enum_value.value, str) else ""
+
+        return (
+            llm_type_lower == enum_name_lower or
+            llm_type_lower == enum_value_lower or
+            llm_type_lower.startswith(enum_name_lower) or # Usecase for snowflake which its type is the provider name + the model name
+            llm_type_lower.startswith(enum_value_lower)
+        )
+
+    return False
+  
+
+def is_support_temperature(llm_type: str, llm_model: str) -> bool:
+    """
+    Check if the LLM model supports temperature setting.
+    """
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    json_file_path = os.path.join(current_dir, 'temperature_supported_models.json')
+
+    try:
+        with open(json_file_path, 'r') as f:
+            temperature_supported_models = json.load(f)
+        
+        # Check if llm_type exists and llm_model is in its list
+        if llm_type in temperature_supported_models:
+            return llm_model in temperature_supported_models[llm_type]
+        
+        return False
+    except (FileNotFoundError, json.JSONDecodeError, KeyError):
+        return False