versionhq 1.1.10.8__py3-none-any.whl → 1.1.11.0__py3-none-any.whl

versionhq/__init__.py

@@ -18,7 +18,7 @@ from versionhq.tool.model import Tool
 from versionhq.tool.composio_tool import ComposioHandler
 
 
-__version__ = "1.1.10.8"
+__version__ = "1.1.11.0"
 __all__ = [
     "Agent",
     "Customer",

versionhq/_utils/vars.py

@@ -0,0 +1,2 @@
+KNOWLEDGE_DIRECTORY="knowledge"
+MAX_FILE_NAME_LENGTH=255

versionhq/agent/TEMPLATES/Backstory.py

@@ -1,4 +1,4 @@
-BACKSTORY_FULL="""You are an expert {role} with deep understanding of {knowledge} and highly skilled in {skillsets}. You have abilities to call the RAG tools that can {rag_tool_overview}. Your primary goal is to identify competitive solutions by leveraging your knowledge and skillsets to achieve the following goal: {goal}."""
+BACKSTORY_FULL="""You are an expert {role} highly skilled in {skills}. You have abilities to query relevant information from the given knowledge sources and use tools such as {tools}. Leveraging these, you will identify competitive solutions to achieve the following goal: {goal}."""
 
 
-BACKSTORY_SHORT="""You are an expert {role} with the right skillsets and knowledge. Your primary goal is to identify competitive solutions by leveraging your knowledge and skillsets to achieve the following goal: {goal}."""
+BACKSTORY_SHORT="""You are an expert {role} with relevant skillsets and abilities to query relevant information from the given knowledge sources. Leveraging these, you will identify competitive solutions to achieve the following goal: {goal}."""

versionhq/agent/default_agents.py

@@ -0,0 +1,15 @@
+from versionhq.agent.model import Agent
+from versionhq.llm.model import DEFAULT_MODEL_NAME
+
+"""
+List up agents to be called across the project.
+"""
+
+client_manager = Agent(role="Client Manager", goal="communicate with clients on the task progress", llm=DEFAULT_MODEL_NAME)
+
+task_evaluator = Agent(
+    role="Task Evaluator",
+    goal="score the output according to the given evaluation criteria.",
+    llm=DEFAULT_MODEL_NAME,
+    llm_config=dict(top_p=0.8, top_k=30, max_tokens=5000, temperature=0.9)
+)

versionhq/agent/model.py

@@ -1,5 +1,6 @@
 import os
 import uuid
+import datetime
 from typing import Any, Dict, List, Optional, TypeVar, Callable, Type
 from typing_extensions import Self
 from dotenv import load_dotenv
@@ -10,8 +11,11 @@ from pydantic_core import PydanticCustomError
 
 from versionhq.llm.model import LLM, DEFAULT_CONTEXT_WINDOW_SIZE, DEFAULT_MODEL_NAME
 from versionhq.tool.model import Tool, ToolSet
+from versionhq.knowledge.model import BaseKnowledgeSource, Knowledge
+from versionhq.memory.contextual_memory import ContextualMemory
+from versionhq.memory.model import ShortTermMemory, LongTermMemory, UserMemory
 from versionhq._utils.logger import Logger
-from versionhq._utils.rpm_controller import RPMController
+from versionhq.agent.rpm_controller import RPMController
 from versionhq._utils.usage_metrics import UsageMetrics
 from versionhq._utils.process_config import process_config
 
@@ -20,9 +24,6 @@ load_dotenv(override=True)
 T = TypeVar("T", bound="Agent")
 
 
-# def _format_answer(agent, answer: str) -> AgentAction | AgentFinish:
-#     return AgentParser(agent=agent).parse(answer)
-
 # def mock_agent_ops_provider():
 #     def track_agent(*args, **kwargs):
 #         def noop(f):
@@ -94,9 +95,20 @@ class Agent(BaseModel):
     role: str = Field(description="role of the agent - used in summary and logs")
     goal: str = Field(description="concise goal of the agent (details are set in the Task instance)")
     backstory: Optional[str] = Field(default=None, description="developer prompt to the llm")
-    knowledge: Optional[str] = Field(default=None, description="external knowledge fed to the agent")
     skillsets: Optional[List[str]] = Field(default_factory=list)
-    tools: Optional[List[Tool | ToolSet | Type[Tool]]] = Field(default_factory=list)
+    tools: Optional[List[InstanceOf[Tool | ToolSet] | Type[Tool] | Any]] = Field(default_factory=list)
+
+    # knowledge
+    knowledge_sources: Optional[List[BaseKnowledgeSource]] = Field(default=None)
+    _knowledge: Optional[Knowledge] = PrivateAttr(default=None)
+
+    # memory
+    use_memory: bool = Field(default=False, description="whether to store/use memory when executing the task")
+    memory_config: Optional[Dict[str, Any]] = Field(default=None, description="configuration for the memory. need to store user_id for UserMemory")
+    short_term_memory: Optional[InstanceOf[ShortTermMemory]] = Field(default=None)
+    long_term_memory: Optional[InstanceOf[LongTermMemory]] = Field(default=None)
+    user_memory: Optional[InstanceOf[UserMemory]] = Field(default=None)
+    embedder_config: Optional[Dict[str, Any]] = Field(default=None, description="embedder configuration for the agent's knowledge")
 
     # prompting
     use_developer_prompt: Optional[bool] = Field(default=True, description="Use developer prompt when calling the llm")
@@ -116,10 +128,10 @@ class Agent(BaseModel):
     respect_context_window: bool = Field(default=True,description="Keep messages under the context window size by summarizing content")
     max_tokens: Optional[int] = Field(default=None, description="max. number of tokens for the agent's execution")
     max_execution_time: Optional[int] = Field(default=None, description="max. execution time for an agent to execute a task")
-    max_rpm: Optional[int] = Field(default=None, description="max. number of requests per minute for the agent execution")
+    max_rpm: Optional[int] = Field(default=None, description="max. number of requests per minute")
     llm_config: Optional[Dict[str, Any]] = Field(default=None, description="other llm config cascaded to the model")
 
-    # config, cache, error handling
+    # cache, error, ops handling
     formatting_errors: int = Field(default=0, description="number of formatting errors.")
     agent_ops_agent_name: str = None
     agent_ops_agent_id: str = None
@@ -132,15 +144,6 @@ class Agent(BaseModel):
             raise PydanticCustomError("may_not_set_field", "This field is not to be set by the user.", {})
 
 
-    # @field_validator(mode="before")
-    # def set_up_from_config(cls) -> None:
-    #     if cls.config is not None:
-    #         try:
-    #             for k, v in cls.config.items():
-    #                 setattr(cls, k, v)
-    #         except:
-    #             pass
-
     @model_validator(mode="before")
     @classmethod
     def process_model_config(cls, values: Dict[str, Any]) -> None:
@@ -165,7 +168,6 @@ class Agent(BaseModel):
         """
 
         self.agent_ops_agent_name = self.role
-        # unaccepted_attributes = ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION_NAME"]
 
         if isinstance(self.llm, LLM):
             llm = self._set_llm_params(self.llm)
@@ -314,29 +316,82 @@ class Agent(BaseModel):
         if self.backstory is None:
             from versionhq.agent.TEMPLATES.Backstory import BACKSTORY_FULL, BACKSTORY_SHORT
             backstory = ""
+            skills = ", ".join([item for item in self.skillsets]) if self.skillsets else ""
+            tools = ", ".join([item.name for item in self.tools if hasattr(item, "name")]) if self.tools else ""
+            role = self.role.lower()
+            goal = self.goal.lower()
 
-            if self.tools or self.knowledge or self.skillsets:
-                backstory = BACKSTORY_FULL.format(
-                    role=self.role,
-                    goal=self.goal,
-                    knowledge=self.knowledge if isinstance(self.knowledge, str) else None,
-                    skillsets=", ".join([item for item in self.skillsets]),
-                    rag_tool_overview=", ".join([item.name for item in self.tools if hasattr(item, "name")]) if self.tools else "",
-                )
+            if self.tools or self.skillsets:
+                backstory = BACKSTORY_FULL.format(role=role, goal=goal, skills=skills, tools=tools)
             else:
-                backstory = BACKSTORY_SHORT.format(role=self.role, goal=self.goal)
+                backstory = BACKSTORY_SHORT.format(role=role, goal=goal)
 
             self.backstory = backstory
 
         return self
 
 
+    @model_validator(mode="after")
+    def set_up_rpm(self) -> Self:
+        """
+        Set up RPM controller.
+        """
+        if self.max_rpm:
+            self._rpm_controller = RPMController(max_rpm=self.max_rpm, _current_rpm=0)
+
+        return self
+
+
+    @model_validator(mode="after")
+    def set_up_knowledge(self) -> Self:
+        if self.knowledge_sources:
+            knowledge_agent_name = f"{self.role.replace(' ', '_')}"
+
+            if isinstance(self.knowledge_sources, list) and all(isinstance(k, BaseKnowledgeSource) for k in self.knowledge_sources):
+                self._knowledge = Knowledge(
+                    sources=self.knowledge_sources,
+                    embedder_config=self.embedder_config,
+                    collection_name=knowledge_agent_name,
+                )
+
+        return self
+
+
+    @model_validator(mode="after")
+    def set_up_memory(self) -> Self:
+        """
+        Set up memories: stm, um
+        """
+
+        if self.use_memory == True:
+            self.long_term_memory = self.long_term_memory if self.long_term_memory else LongTermMemory()
+            self.short_term_memory = self.short_term_memory if self.short_term_memory else ShortTermMemory(agent=self, embedder_config=self.embedder_config)
+
+            if hasattr(self, "memory_config") and self.memory_config is not None:
+                user_id = self.memory_config.get("user_id", None)
+                if user_id:
+                    self.user_memory = self.user_memory if self.user_memory else UserMemory(agent=self, user_id=user_id)
+            else:
+                self.user_memory = None
+
+        return self
+
+
+    def _train(self) -> Self:
+        """
+        Fine-tuned the base model using OpenAI train framework.
+        """
+        if not isinstance(self.llm, LLM):
+            pass
+
+
     def invoke(
         self,
         prompts: str,
         response_format: Optional[Dict[str, Any]] = None,
-        tools: Optional[List[Tool | ToolSet | Type[Tool]]] = None,
-        tool_res_as_final: bool = False
+        tools: Optional[List[InstanceOf[Tool]| InstanceOf[ToolSet] | Type[Tool]]] = None,
+        tool_res_as_final: bool = False,
+        task: Any = None
         ) -> Dict[str, Any]:
         """
         Create formatted prompts using the developer prompt and the agent's backstory, then call the base model.
@@ -347,40 +402,51 @@ class Agent(BaseModel):
         task_execution_counter = 0
         iterations = 0
         raw_response = None
-        messages = []
 
-        messages.append({"role": "user", "content": prompts})
+        messages = []
+        messages.append({ "role": "user", "content": prompts })
         if self.use_developer_prompt:
-            messages.append({"role": "system", "content": self.backstory})
-        self._logger.log(level="info", message=f"Messages sent to the model: {messages}", color="blue")
+            messages.append({ "role": "system", "content": self.backstory })
 
         try:
-            if tool_res_as_final is True:
+            if self._rpm_controller and self.max_rpm:
+                self._rpm_controller.check_or_wait()
+
+            self._logger.log(level="info", message=f"Messages sent to the model: {messages}", color="blue")
+
+            if tool_res_as_final:
                 func_llm = self.function_calling_llm if self.function_calling_llm and self.function_calling_llm._supports_function_calling() else LLM(model=DEFAULT_MODEL_NAME)
                 raw_response = func_llm.call(messages=messages, tools=tools, tool_res_as_final=True)
+                task.tokens = func_llm._tokens
             else:
                 raw_response = self.llm.call(messages=messages, response_format=response_format, tools=tools)
+                task.tokens = self.llm._tokens
 
             task_execution_counter += 1
             self._logger.log(level="info", message=f"Agent response: {raw_response}", color="blue")
-
+            return raw_response
 
         except Exception as e:
             self._logger.log(level="error", message=f"An error occured. The agent will retry: {str(e)}", color="red")
 
-            while not raw_response and task_execution_counter < self.max_retry_limit:
-                while not raw_response and iterations < self.maxit:
+            while not raw_response and task_execution_counter <= self.max_retry_limit:
+                while (not raw_response or raw_response == "" or raw_response is None) and iterations < self.maxit:
+                    if self.max_rpm and self._rpm_controller:
+                        self._rpm_controller.check_or_wait()
+
                     raw_response = self.llm.call(messages=messages, response_format=response_format, tools=tools)
+                    task.tokens = self.llm._tokens
                     iterations += 1
 
                 task_execution_counter += 1
                 self._logger.log(level="info", message=f"Agent #{task_execution_counter} response: {raw_response}", color="blue")
+                return raw_response
 
             if not raw_response:
                 self._logger.log(level="error", message="Received None or empty response from the model", color="red")
                 raise ValueError("Invalid response from LLM call - None or empty.")
 
-        return raw_response
+
 
 
     def execute_task(self, task, context: Optional[str] = None, task_tools: Optional[List[Tool | ToolSet]] = list()) -> str:
@@ -390,14 +456,35 @@ class Agent(BaseModel):
         The agent must consider the context to excute the task as well when it is given.
         """
         from versionhq.task.model import Task
+        from versionhq.knowledge._utils import extract_knowledge_context
 
         task: InstanceOf[Task] = task
-        tools: Optional[List[Tool | ToolSet | Type[Tool]]] = task_tools + self.tools if task.can_use_agent_tools else task_tools
+        tools: Optional[List[InstanceOf[Tool]| InstanceOf[ToolSet] | Type[Tool]]] = task_tools + self.tools if task.can_use_agent_tools else task_tools
+
+        if self.max_rpm and self._rpm_controller:
+            self._rpm_controller._reset_request_count()
 
         task_prompt = task.prompt(model_provider=self.llm.provider)
         if context is not task.prompt_context:
             task_prompt += context
 
+        if self._knowledge:
+            agent_knowledge = self._knowledge.query(query=[task_prompt,])
+            if agent_knowledge:
+                agent_knowledge_context = extract_knowledge_context(knowledge_snippets=agent_knowledge)
+                if agent_knowledge_context:
+                    task_prompt += agent_knowledge_context
+
+
+        if self.use_memory == True:
+            contextual_memory = ContextualMemory(
+                memory_config=self.memory_config, stm=self.short_term_memory, ltm=self.long_term_memory, um=self.user_memory
+            )
+            memory = contextual_memory.build_context_for_task(task=task, context=context)
+            if memory.strip() != "":
+                task_prompt += memory.strip()
+
+
         # if self.team and self.team._train:
         #     task_prompt = self._training_handler(task_prompt=task_prompt)
         # else:
@@ -410,6 +497,7 @@ class Agent(BaseModel):
                 response_format=task._structure_response_format(model_provider=self.llm.provider),
                 tools=tools,
                 tool_res_as_final=task.tool_res_as_final,
+                task=task
             )
 
         except Exception as e:

versionhq/agent/parser.py

@@ -44,26 +44,6 @@ class OutputParserException(Exception):
 
 
 class AgentParser:
-    """
-    Parses ReAct-style LLM calls that have a single tool input.
-
-    Expects output to be in one of two formats.
-
-    If the output signals that an action should be taken,
-    should be in the below format. This will result in an AgentAction
-    being returned.
-
-    Thought: agent thought here
-    Action: search
-    Action Input: what is the temperature in SF?
-
-    If the output signals that a final answer should be given,
-    should be in the below format. This will result in an AgentFinish
-    being returned.
-
-    Thought: agent thought here
-    Final Answer: The temperature is 100 degrees
-    """
 
     agent: Any = None
 
@@ -111,6 +91,7 @@ class AgentParser:
             # self.agent.increment_formatting_errors()
             raise OutputParserException(error)
 
+
     def _extract_thought(self, text: str) -> str:
         regex = r"(.*?)(?:\n\nAction|\n\nFinal Answer)"
         thought_match = re.search(regex, text, re.DOTALL)
@@ -118,10 +99,12 @@ class AgentParser:
             return thought_match.group(1).strip()
         return ""
 
+
     def _clean_action(self, text: str) -> str:
         """Clean action string by removing non-essential formatting characters."""
         return re.sub(r"^\s*\*+\s*|\s*\*+\s*$", "", text).strip()
 
+
     def _safe_repair_json(self, tool_input: str) -> str:
         UNABLE_TO_REPAIR_JSON_RESULTS = ['""', "{}"]
 

versionhq/{_utils → agent}/rpm_controller.py

@@ -10,7 +10,7 @@ from versionhq._utils.logger import Logger
 
 class RPMController(BaseModel):
     max_rpm: Optional[int] = Field(default=None)
-    logger: Logger = Field(default_factory=lambda: Logger(verbose=False))
+    _logger: Logger = PrivateAttr(default_factory=lambda: Logger(verbose=True))
     _current_rpm: int = PrivateAttr(default=0)
     _timer: Optional[threading.Timer] = PrivateAttr(default=None)
     _lock: Optional[threading.Lock] = PrivateAttr(default=None)
@@ -24,38 +24,45 @@ class RPMController(BaseModel):
                 self._reset_request_count()
         return self
 
-    def check_or_wait(self) -> bool:
+
+    def _check_and_increment(self) -> bool:
         if self.max_rpm is None:
             return True
 
-        def _check_and_increment():
-            if self.max_rpm is not None and self._current_rpm < self.max_rpm:
-                self._current_rpm += 1
-                return True
-            elif self.max_rpm is not None:
-                self.logger.log(
-                    "info", "Max RPM reached, waiting for next minute to start."
-                )
-                self._wait_for_next_minute()
-                self._current_rpm = 1
-                return True
+        elif self.max_rpm is not None and self._current_rpm < self.max_rpm:
+            self._current_rpm += 1
+            return True
+
+        elif self.max_rpm is not None and self._current_rpm >= self.max_rpm:
+            self._logger.log(level="info", message="Max RPM reached, waiting for next minute to start.", color="yellow")
+            self._wait_for_next_minute()
+            self._current_rpm = 1 # restart
             return True
 
+        else:
+            return False
+
+    def check_or_wait(self) -> bool:
         if self._lock:
             with self._lock:
-                return _check_and_increment()
+                return self._check_and_increment()
         else:
-            return _check_and_increment()
+            return self._check_and_increment()
+
+        return False
+
 
     def stop_rpm_counter(self) -> None:
         if self._timer:
             self._timer.cancel()
             self._timer = None
 
+
     def _wait_for_next_minute(self) -> None:
         time.sleep(60)
         self._current_rpm = 0
 
+
     def _reset_request_count(self) -> None:
         def _reset():
             self._current_rpm = 0

versionhq/knowledge/_utils.py

@@ -0,0 +1,11 @@
+from typing import List, Dict, Any
+
+
+def extract_knowledge_context(knowledge_snippets: List[Dict[str, Any]]) -> str:
+    """
+    Extract knowledge from the task prompt.
+    """
+
+    valid_snippets = [result["context"] for result in knowledge_snippets if result and result.get("context")]
+    snippet = "\n".join(valid_snippets)
+    return f"Additional Information: {snippet}" if valid_snippets else ""

versionhq/knowledge/embedding.py

@@ -0,0 +1,192 @@
+import os
+from typing import Any, Dict, cast
+
+from chromadb import Documents, EmbeddingFunction, Embeddings
+from chromadb.api.types import validate_embedding_function
+
+from versionhq._utils.logger import Logger
+
+
+class EmbeddingConfigurator:
+    def __init__(self):
+        self.embedding_functions = {
+            "openai": self._configure_openai,
+            "azure": self._configure_azure,
+            "ollama": self._configure_ollama,
+            "vertexai": self._configure_vertexai,
+            "google": self._configure_google,
+            "cohere": self._configure_cohere,
+            "voyageai": self._configure_voyageai,
+            "bedrock": self._configure_bedrock,
+            "huggingface": self._configure_huggingface,
+            "watson": self._configure_watson,
+        }
+
+    def configure_embedder(self, embedder_config: Dict[str, Any] | None = None) -> EmbeddingFunction:
+        """
+        Configures and returns an embedding function based on the provided config.
+        """
+
+        if embedder_config is None:
+            return self._create_default_embedding_function()
+
+        provider = embedder_config.get("provider")
+        config = embedder_config.get("config", {})
+        model_name = config.get("model")
+
+        if isinstance(provider, EmbeddingFunction):
+            try:
+                validate_embedding_function(provider)
+                return provider
+            except Exception as e:
+                raise ValueError(f"Invalid custom embedding function: {str(e)}")
+
+        if provider not in self.embedding_functions:
+            raise Exception(
+                f"Unsupported embedding provider: {provider}, supported providers: {list(self.embedding_functions.keys())}"
+            )
+
+        return self.embedding_functions[provider](config, model_name)
+
+    @staticmethod
+    def _create_default_embedding_function():
+        from chromadb.utils.embedding_functions.openai_embedding_function import (
+            OpenAIEmbeddingFunction,
+        )
+
+        return OpenAIEmbeddingFunction(
+            api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small"
+        )
+
+    @staticmethod
+    def _configure_openai(config, model_name):
+        from chromadb.utils.embedding_functions.openai_embedding_function import (
+            OpenAIEmbeddingFunction,
+        )
+
+        return OpenAIEmbeddingFunction(
+            api_key=config.get("api_key") or os.getenv("OPENAI_API_KEY"),
+            model_name=model_name,
+        )
+
+    @staticmethod
+    def _configure_azure(config, model_name):
+        from chromadb.utils.embedding_functions.openai_embedding_function import (
+            OpenAIEmbeddingFunction,
+        )
+
+        return OpenAIEmbeddingFunction(
+            api_key=config.get("api_key"),
+            api_base=config.get("api_base"),
+            api_type=config.get("api_type", "azure"),
+            api_version=config.get("api_version"),
+            model_name=model_name,
+        )
+
+    @staticmethod
+    def _configure_ollama(config, model_name):
+        from chromadb.utils.embedding_functions.ollama_embedding_function import (
+            OllamaEmbeddingFunction,
+        )
+
+        return OllamaEmbeddingFunction(
+            url=config.get("url", "http://localhost:11434/api/embeddings"),
+            model_name=model_name,
+        )
+
+    @staticmethod
+    def _configure_vertexai(config, model_name):
+        from chromadb.utils.embedding_functions.google_embedding_function import (
+            GoogleVertexEmbeddingFunction,
+        )
+
+        return GoogleVertexEmbeddingFunction(model_name=model_name, api_key=config.get("api_key"))
+
+    @staticmethod
+    def _configure_google(config, model_name):
+        from chromadb.utils.embedding_functions.google_embedding_function import (
+            GoogleGenerativeAiEmbeddingFunction,
+        )
+
+        return GoogleGenerativeAiEmbeddingFunction(model_name=model_name, api_key=config.get("api_key"))
+
+    @staticmethod
+    def _configure_cohere(config, model_name):
+        from chromadb.utils.embedding_functions.cohere_embedding_function import (
+            CohereEmbeddingFunction,
+        )
+
+        return CohereEmbeddingFunction(
+            model_name=model_name,
+            api_key=config.get("api_key"),
+        )
+
+    @staticmethod
+    def _configure_voyageai(config, model_name):
+        from chromadb.utils.embedding_functions.voyageai_embedding_function import (
+            VoyageAIEmbeddingFunction,
+        )
+
+        return VoyageAIEmbeddingFunction(
+            model_name=model_name,
+            api_key=config.get("api_key"),
+        )
+
+    @staticmethod
+    def _configure_bedrock(config, model_name):
+        from chromadb.utils.embedding_functions.amazon_bedrock_embedding_function import (
+            AmazonBedrockEmbeddingFunction,
+        )
+
+        return AmazonBedrockEmbeddingFunction(
+            session=config.get("session"),
+        )
+
+    @staticmethod
+    def _configure_huggingface(config, model_name):
+        from chromadb.utils.embedding_functions.huggingface_embedding_function import (
+            HuggingFaceEmbeddingServer,
+        )
+
+        return HuggingFaceEmbeddingServer(
+            url=config.get("api_url"),
+        )
+
+    @staticmethod
+    def _configure_watson(config, model_name):
+        try:
+            import ibm_watsonx_ai.foundation_models as watson_models
+            from ibm_watsonx_ai import Credentials
+            from ibm_watsonx_ai.metanames import EmbedTextParamsMetaNames as EmbedParams
+        except ImportError as e:
+            raise ImportError(
+                "IBM Watson dependencies are not installed. Please install them to use Watson embedding."
+            ) from e
+
+        class WatsonEmbeddingFunction(EmbeddingFunction):
+            def __call__(self, input: Documents) -> Embeddings:
+                if isinstance(input, str):
+                    input = [input]
+
+                embed_params = {
+                    EmbedParams.TRUNCATE_INPUT_TOKENS: 3,
+                    EmbedParams.RETURN_OPTIONS: {"input_text": True},
+                }
+
+                embedding = watson_models.Embeddings(
+                    model_id=config.get("model"),
+                    params=embed_params,
+                    credentials=Credentials(
+                        api_key=config.get("api_key"), url=config.get("api_url")
+                    ),
+                    project_id=config.get("project_id"),
+                )
+
+                try:
+                    embeddings = embedding.embed_documents(input)
+                    return cast(Embeddings, embeddings)
+                except Exception as e:
+                    Logger(verbose=True).log(level="error", message=f"Error during Watson embedding: {str(e)}", color="red")
+                    raise e
+
+        return WatsonEmbeddingFunction()