agentica 0.1.4__tar.gz → 0.1.5__tar.gz

{agentica-0.1.4 → agentica-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agentica
-Version: 0.1.4
+Version: 0.1.5
 Summary: LLM agents
 Home-page: https://github.com/shibing624/agentica
 Author: XuMing
@@ -25,6 +25,7 @@ Requires-Dist: python-dotenv
 Requires-Dist: pydantic
 Requires-Dist: requests
 Requires-Dist: sqlalchemy
+Requires-Dist: scikit-learn
 Requires-Dist: tqdm
 Requires-Dist: rich
 
@@ -91,13 +92,16 @@ pip install .
 
 ## Getting Started
 
-### 运行示例
+#### 1. Install requirements
 
 ```shell
 git clone https://github.com/shibing624/agentica.git
 cd agentica
 pip install -r requirements.txt
+```
 
+#### 2. Run the example
+```shell
 # Copying required .env file, and fill in the LLM api key
 cp .env.example ~/.agentica/.env
 
@@ -112,11 +116,11 @@ python web_search_deepseek_demo.py
 自动调用google搜索工具，示例[examples/web_search_deepseek_demo.py](https://github.com/shibing624/agentica/blob/main/examples/web_search_deepseek_demo.py)
 
 ```python
-from agentica import Assistant, DeepseekLLM
+from agentica import Assistant, Deepseek
 from agentica.tools.search_serper import SearchSerperTool
 
 m = Assistant(
-  llm=DeepseekLLM(),
+  llm=Deepseek(),
   description="You are a helpful ai assistant.",
   show_tool_calls=True,
   # Enable the assistant to search the knowledge base
@@ -198,9 +202,6 @@ bash start.sh
 
 <img src="https://github.com/shibing624/agentica/blob/main/docs/wechat.jpeg" width="200" />
 
-<img src="https://github.com/shibing624/agentica/blob/main/docs/wechat_group.jpg" width="200" />
-
-
 ## Citation
 
 如果你在研究中使用了`agentica`，请按如下格式引用：

{agentica-0.1.4 → agentica-0.1.5}/README.md

@@ -61,13 +61,16 @@ pip install .
 
 ## Getting Started
 
-### 运行示例
+#### 1. Install requirements
 
 ```shell
 git clone https://github.com/shibing624/agentica.git
 cd agentica
 pip install -r requirements.txt
+```
 
+#### 2. Run the example
+```shell
 # Copying required .env file, and fill in the LLM api key
 cp .env.example ~/.agentica/.env
 
@@ -82,11 +85,11 @@ python web_search_deepseek_demo.py
 自动调用google搜索工具，示例[examples/web_search_deepseek_demo.py](https://github.com/shibing624/agentica/blob/main/examples/web_search_deepseek_demo.py)
 
 ```python
-from agentica import Assistant, DeepseekLLM
+from agentica import Assistant, Deepseek
 from agentica.tools.search_serper import SearchSerperTool
 
 m = Assistant(
-  llm=DeepseekLLM(),
+  llm=Deepseek(),
   description="You are a helpful ai assistant.",
   show_tool_calls=True,
   # Enable the assistant to search the knowledge base
@@ -168,9 +171,6 @@ bash start.sh
 
 <img src="https://github.com/shibing624/agentica/blob/main/docs/wechat.jpeg" width="200" />
 
-<img src="https://github.com/shibing624/agentica/blob/main/docs/wechat_group.jpg" width="200" />
-
-
 ## Citation
 
 如果你在研究中使用了`agentica`，请按如下格式引用：

agentica-0.1.5/agentica/__init__.py

@@ -0,0 +1,66 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing(xuming624@qq.com)
+@description: 
+"""
+from agentica.version import __version__  # noqa, isort: skip
+from agentica.config import AGENTICA_HOME, AGENTICA_DOTENV_PATH, SMART_LLM, FAST_LLM  # noqa, isort: skip
+
+# llm
+from agentica.llm.openai_chat import OpenAIChat
+from agentica.llm.azure_openai import AzureOpenAIChat
+from agentica.llm.together import Together
+from agentica.llm.deepseek import Deepseek
+from agentica.llm.moonshot import Moonshot
+from agentica.llm.ollama_chat import OllamaChat
+from agentica.llm.ollama_tools import OllamaTools
+from agentica.llm.claude import Claude
+
+# memory
+from agentica.message import Message
+from agentica.memory import (
+    Memory,
+    MemoryRow,
+    MemoryDb,
+    CsvMemoryDb,
+    InMemoryDb,
+    MemoryRetrieval,
+    AssistantMemory,
+    MemoryClassifier,
+    MemoryManager
+)
+from agentica.template import PromptTemplate
+# rag
+from agentica.knowledge.knowledge_base import KnowledgeBase
+from agentica.references import References
+from agentica.run_record import RunRecord
+from agentica.document import Document
+# vectordb
+from agentica.vectordb.base import VectorDb
+from agentica.vectordb.memory_vectordb import MemoryVectorDb
+# emb
+from agentica.emb.base import Emb
+from agentica.emb.openai import OpenAIEmb
+from agentica.emb.azure_openai import AzureOpenAIEmb
+from agentica.emb.hash_emb import HashEmb
+from agentica.emb.together import TogetherEmb
+from agentica.emb.fireworks import FireworksEmb
+from agentica.emb.text2vec import Text2VecEmb
+from agentica.emb.word2vec import Word2VecEmb
+
+# file
+from agentica.file.base import File
+from agentica.file.csv import CsvFile
+from agentica.file.txt import TextFile
+
+# storage
+from agentica.storage.base import AssistantStorage
+from agentica.storage.pg_storage import PgAssistantStorage
+from agentica.storage.sqlite_storage import SqlAssistantStorage
+# tool
+from agentica.tools.base import Tool, Toolkit, Function, FunctionCall
+# assistant
+from agentica.assistant import Assistant
+from agentica.python_assistant import PythonAssistant
+from agentica.task import Task
+from agentica.workflow import Workflow

{agentica-0.1.4 → agentica-0.1.5}/agentica/assistant.py

@@ -25,21 +25,22 @@ from typing import (
     AsyncIterator,
 )
 from uuid import uuid4
-
+from pathlib import Path
 from pydantic import BaseModel, ConfigDict, field_validator, ValidationError
 
 from agentica.document import Document
 from agentica.knowledge.knowledge_base import KnowledgeBase
 from agentica.llm.base import LLM
-from agentica.llm.openai_llm import OpenAILLM
+from agentica.llm.openai_chat import OpenAIChat
 from agentica.memory import AssistantMemory, Memory
 from agentica.message import Message
 from agentica.references import References
 from agentica.run_record import RunRecord
-from agentica.sqlite_storage import SqliteStorage
-from agentica.tool import Tool, Toolkit, Function
+from agentica.storage.base import AssistantStorage
+from agentica.tools.base import Tool, Toolkit, Function
 from agentica.utils.log import logger, set_log_level_to_debug, print_llm_stream
 from agentica.utils.misc import merge_dictionaries, remove_indent
+from agentica.template import PromptTemplate
 from agentica.utils.timer import Timer
 
 
@@ -80,6 +81,8 @@ class Assistant(BaseModel):
     create_memories: bool = False
     # Update memory after each run
     update_memory_after_run: bool = True
+    # Force update memory after each run
+    force_update_memory_after_run: bool = False
 
     # -*- Assistant Knowledge Base
     knowledge_base: Optional[KnowledgeBase] = None
@@ -87,7 +90,7 @@ class Assistant(BaseModel):
     add_references_to_prompt: bool = False
 
     # -*- Assistant Storage
-    storage: Optional[SqliteStorage] = None
+    storage: Optional[AssistantStorage] = None
     # RunRecord from the database: DO NOT SET MANUALLY
     db_row: Optional[RunRecord] = None
     # -*- Assistant Tools
@@ -130,6 +133,8 @@ class Assistant(BaseModel):
     #
     # -*- System prompt: provide the system prompt as a string
     system_prompt: Optional[str] = None
+    # -*- System prompt template: provide the system prompt as a PromptTemplate
+    system_prompt_template: Optional[PromptTemplate] = None
     # If True, build a default system prompt using instructions and extra_instructions
     build_default_system_prompt: bool = True
     # -*- Settings for building the default system prompt
@@ -147,7 +152,7 @@ class Assistant(BaseModel):
     add_to_system_prompt: Optional[str] = None
     # If True, add instructions for using the knowledge base to the system prompt if knowledge base is provided
     add_knowledge_base_instructions: bool = True
-    # If True, add instructions to return "I dont know" when the assistant does not know the answer.
+    # If True, add instructions to return "I don't know" when the assistant does not know the answer.
     prevent_hallucinations: bool = False
     # If True, add instructions to prevent prompt injection attacks
     prevent_prompt_injection: bool = False
@@ -162,6 +167,8 @@ class Assistant(BaseModel):
     # -*- User prompt: provide the user prompt as a string
     # Note: this will ignore the message sent to the run function
     user_prompt: Optional[Union[List, Dict, str]] = None
+    # -*- User prompt template: provide the user prompt as a PromptTemplate
+    user_prompt_template: Optional[PromptTemplate] = None
     # If True, build a default user prompt using references and chat history
     build_default_user_prompt: bool = True
     # Function to get references for the user_prompt
@@ -187,6 +194,8 @@ class Assistant(BaseModel):
     output: Optional[Any] = None
     # Save the output to output_dir
     output_dir: Optional[str] = "outputs"
+    # Save llm messages to output_dir
+    save_llm_messages: bool = False
     # Save the output to a file with this name, if provided save the output to a file
     output_file_name: Optional[str] = None
 
@@ -270,8 +279,10 @@ class Assistant(BaseModel):
     def update_llm(self) -> None:
         if self.llm is None:
             logger.debug("LLM not set. Using OpenAILLM")
-            self.llm = OpenAILLM()
-        logger.debug(f"Using LLM: {self.llm}")
+            self.llm = OpenAIChat()
+            logger.info(f"Using LLM: {self.llm}")
+        else:
+            logger.debug(f"Using LLM: {self.llm}")
 
         # Set response_format if it is not set on the llm
         if self.output_model is not None and self.llm.response_format is None:
@@ -289,7 +300,10 @@ class Assistant(BaseModel):
                 self.llm.add_tool(self.get_tool_call_history)
             if self.create_memories:
                 self.llm.add_tool(self.update_memory)
+        if self.force_update_memory_after_run:
+            self.update_memory_after_run = True
         if self.knowledge_base is not None:
+            self.add_references_to_prompt = True
             if self.search_knowledge:
                 self.llm.add_tool(self.search_knowledge_base)
             if self.update_knowledge:
@@ -555,7 +569,13 @@ class Assistant(BaseModel):
                 sys_prompt += f"\n{self.get_json_output_prompt()}"
                 return sys_prompt
             return self.system_prompt
-
+        # If the system_prompt_template is set, build the system_prompt using the template
+        if self.system_prompt_template is not None:
+            system_prompt_kwargs = {"assistant": self}
+            system_prompt_from_template = self.system_prompt_template.get_prompt(**system_prompt_kwargs)
+            if system_prompt_from_template is not None and self.output_model is not None:
+                system_prompt_from_template += f"\n{self.get_json_output_prompt()}"
+            return system_prompt_from_template
         # If build_default_system_prompt is False, return None
         if not self.build_default_system_prompt:
             return None
@@ -564,7 +584,7 @@ class Assistant(BaseModel):
             raise Exception("LLM not set")
 
         # -*- Build a list of instructions for the Assistant
-        instructions = self.instructions.copy() if self.instructions is not None else []
+        instructions = self.instructions.copy() if self.instructions is not None else None
         # Add default instructions
         if instructions is None:
             instructions = []
@@ -716,12 +736,12 @@ class Assistant(BaseModel):
 
         if self.chat_history_function is not None:
             chat_history_kwargs = {"conversation": self}
-            return self.chat_history_function(**chat_history_kwargs)
+            return remove_indent(self.chat_history_function(**chat_history_kwargs))
 
         formatted_history = self.memory.get_formatted_chat_history(num_messages=self.num_history_messages)
         if formatted_history == "":
             return None
-        return formatted_history
+        return remove_indent(formatted_history)
 
     def get_user_prompt(
             self,
@@ -735,6 +755,16 @@ class Assistant(BaseModel):
         # Note: this ignores the message provided to the run function
         if self.user_prompt is not None:
             return self.user_prompt
+        # If the user_prompt_template is set, return the user_prompt from the template
+        if self.user_prompt_template is not None:
+            user_prompt_kwargs = {
+                "assistant": self,
+                "message": message,
+                "references": references,
+                "chat_history": chat_history,
+            }
+            _user_prompt_from_template = self.user_prompt_template.get_prompt(**user_prompt_kwargs)
+            return _user_prompt_from_template
 
         if message is None:
             return None
@@ -875,8 +905,12 @@ class Assistant(BaseModel):
         if user_message is not None:
             self.memory.add_chat_message(message=user_message)
             # Update the memory with the user message if needed
+            memory_content = f"user:{user_message.get_content_string()}\nassistant:{llm_response}"
             if self.create_memories and self.update_memory_after_run:
-                self.memory.update_memory(input_text=user_message.get_content_string())
+                if self.force_update_memory_after_run:
+                    self.memory.update_memory(input_text=memory_content, force=True)
+                else:
+                    self.memory.update_memory(input_text=memory_content)
 
         # Build the LLM response message to add to the memory - this is added to the chat_history
         llm_response_message = Message(role="assistant", content=llm_response)
@@ -895,24 +929,28 @@ class Assistant(BaseModel):
         self.write_to_storage()
 
         # Save llm_messages to file
-        try:
-            os.makedirs(self.output_dir, exist_ok=True)
-            save_file = os.path.join(self.output_dir, f"output_{self.run_id}.json")
-            messages_str = json.dumps(
-                [i.dict() for i in self.memory.llm_messages], indent=2, ensure_ascii=False
-            )
-            with open(save_file, "w", encoding='utf-8') as f:
-                f.write(messages_str)
-            logger.info(f"Saved messages to file: {save_file}")
-        except Exception as e:
-            logger.warning(f"Failed to save output to file: {e}")
+        if self.save_llm_messages:
+            try:
+                os.makedirs(self.output_dir, exist_ok=True)
+                save_file = os.path.join(self.output_dir, f"output_{self.run_id}.json")
+                messages_str = json.dumps(
+                    [i.dict() for i in self.memory.llm_messages], indent=2, ensure_ascii=False
+                )
+                with open(save_file, "w", encoding='utf-8') as f:
+                    f.write(messages_str)
+                logger.info(f"Saved messages to file: {save_file}")
+            except Exception as e:
+                logger.warning(f"Failed to save output to file: {e}")
 
         # Save output_file_name file
         if self.output_file_name:
             try:
+                os.makedirs(self.output_dir, exist_ok=True)
                 save_file = os.path.join(self.output_dir, self.output_file_name)
-                with open(save_file, "w", encoding='utf-8') as f:
-                    f.write(self.output)
+                fn_path = Path(save_file)
+                if not fn_path.parent.exists():
+                    fn_path.parent.mkdir(parents=True, exist_ok=True)
+                fn_path.write_text(self.output)
                 logger.info(f"Saved output to file: {save_file}")
             except Exception as e:
                 logger.warning(f"Failed to save output to file: {e}")
@@ -1064,8 +1102,12 @@ class Assistant(BaseModel):
         if user_message is not None:
             self.memory.add_chat_message(message=user_message)
             # Update the memory with the user message if needed
+            memory_content = f"user:{user_message.get_content_string()}\nassistant:{llm_response}"
             if self.update_memory_after_run:
-                self.memory.update_memory(input_text=user_message.get_content_string())
+                if self.force_update_memory_after_run:
+                    self.memory.update_memory(input_text=memory_content, force=True)
+                else:
+                    self.memory.update_memory(input_text=memory_content)
 
         # Build the LLM response message to add to the memory - this is added to the chat_history
         llm_response_message = Message(role="assistant", content=llm_response)

agentica-0.1.5/agentica/config.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing(xuming624@qq.com)
+@description: 
+"""
+import os
+from datetime import datetime
+from dotenv import load_dotenv  # noqa
+from loguru import logger  # noqa, need to import logger here to avoid circular import
+
+AGENTICA_HOME = os.getenv("AGENTICA_HOME", os.path.expanduser("~/.agentica"))
+
+# Load environment variables from .env file
+AGENTICA_DOTENV_PATH = os.getenv("AGENTICA_DOTENV_PATH", f"{AGENTICA_HOME}/.env")
+
+if load_dotenv(AGENTICA_DOTENV_PATH, override=True):
+    logger.info(f"Loaded AGENTICA_DOTENV_PATH: {AGENTICA_DOTENV_PATH}")
+
+AGENTICA_DATA_DIR = os.getenv("AGENTICA_DATA_DIR", f"{AGENTICA_HOME}/data")
+AGENTICA_LOG_LEVEL = os.getenv("AGENTICA_LOG_LEVEL", "INFO")
+AGENTICA_LOG_FILE = os.getenv("AGENTICA_LOG_FILE")
+if AGENTICA_LOG_LEVEL.upper() == "DEBUG":
+    formatted_date = datetime.now().strftime("%Y%m%d")
+    default_log_file = f"{AGENTICA_HOME}/logs/{formatted_date}.log"
+    AGENTICA_LOG_FILE = os.getenv("AGENTICA_LOG_FILE", default_log_file)
+    logger.debug(f"AGENTICA_LOG_LEVEL: DEBUG, AGENTICA_LOG_FILE: {AGENTICA_LOG_FILE}")
+
+SMART_LLM = os.getenv("SMART_LLM")
+FAST_LLM = os.getenv("FAST_LLM")

{agentica-0.1.4 → agentica-0.1.5}/agentica/document.py

@@ -35,6 +35,7 @@ class Document(BaseModel):
             self.embedding, self.usage = _embedder.get_embedding_and_usage(self.content)
         else:
             self.embedding = _embedder.get_embedding(self.content)
+            self.usage = None
 
     def to_dict(self) -> Dict[str, Any]:
         """Returns a dictionary representation of the document"""

agentica-0.1.4/agentica/emb/azure_emb.py → agentica-0.1.5/agentica/emb/azure_openai.py

@@ -5,11 +5,10 @@
 part of the code from https://github.com/phidatahq/phidata
 """
 from os import getenv
-from typing import Optional, Dict, List, Tuple, Any
+from typing import Optional, Dict, List, Tuple, Any, Literal
 
 from openai import AzureOpenAI as AzureOpenAIClient
 from openai.types.create_embedding_response import CreateEmbeddingResponse
-from typing_extensions import Literal
 
 from agentica.emb.base import Emb
 from agentica.utils.log import logger

{agentica-0.1.4 → agentica-0.1.5}/agentica/emb/base.py

@@ -6,12 +6,13 @@ from pydantic import BaseModel, ConfigDict
 class Emb(BaseModel):
     """Base class for managing embedders"""
 
-    dimensions: int = 1536
+    dimensions: Optional[int] = 1536
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     def get_embedding(self, text: str) -> List[float]:
         raise NotImplementedError
 
-    def get_embeddings(self, texts: List[str]) -> List[List[float]]:
-        raise NotImplementedError
+    def get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict]]:
+        embedding = self.get_embedding(text)
+        return embedding, None

agentica-0.1.5/agentica/emb/fireworks.py

@@ -0,0 +1,16 @@
+"""
+@author:XuMing(xuming624@qq.com)
+@description:
+part of the code from https://github.com/phidatahq/phidata
+"""
+from os import getenv
+from typing import Optional
+
+from agentica.emb.openai import OpenAIEmb
+
+
+class FireworksEmb(OpenAIEmb):
+    model: str = "nomic-ai/nomic-embed-text-v1.5"
+    dimensions: int = 768
+    api_key: Optional[str] = getenv("FIREWORKS_API_KEY")
+    base_url: str = "https://api.fireworks.ai/inference/v1"

agentica-0.1.5/agentica/emb/genimi.py

@@ -0,0 +1,60 @@
+"""
+@author:XuMing(xuming624@qq.com)
+@description:
+part of the code from https://github.com/phidatahq/phidata
+"""
+from typing import Optional, Dict, List, Tuple, Any, Union
+
+from agentica.emb.base import Emb
+from agentica.utils.log import logger
+
+try:
+    import google.generativeai as genai
+    from google.generativeai.types.text_types import EmbeddingDict, BatchEmbeddingDict
+except ImportError:
+    raise ImportError("`google-generativeai` not installed. Please install it using `pip install google-generativeai`")
+
+
+class GeminiEmb(Emb):
+    model: str = "models/embedding-001"
+    task_type: str = "RETRIEVAL_QUERY"
+    title: Optional[str] = None
+    dimensions: Optional[int] = None
+    api_key: Optional[str] = None
+    request_params: Optional[Dict[str, Any]] = None
+    client_params: Optional[Dict[str, Any]] = None
+    gemini_client: Optional[genai.embed_content] = None
+
+    @property
+    def client(self):
+        if self.gemini_client:
+            return self.gemini_client
+        _client_params: Dict[str, Any] = {}
+        if self.api_key:
+            _client_params["api_key"] = self.api_key
+        if self.client_params:
+            _client_params.update(self.client_params)
+        self.gemini_client = genai
+        self.gemini_client.configure(**_client_params)
+        return self.gemini_client
+
+    def _response(self, text: str) -> Union[EmbeddingDict, BatchEmbeddingDict]:
+        _request_params: Dict[str, Any] = {
+            "content": text,
+            "model": self.model,
+            "output_dimensionality": self.dimensions,
+            "task_type": self.task_type,
+            "title": self.title,
+        }
+        if self.request_params:
+            _request_params.update(self.request_params)
+        return self.client.embed_content(**_request_params)
+
+    def get_embedding(self, text: str) -> List[float]:
+        response = self._response(text=text)
+        try:
+            return response.get("embedding", [])
+        except Exception as e:
+            logger.warning(e)
+            return []
+

agentica-0.1.5/agentica/emb/huggingface.py

@@ -0,0 +1,53 @@
+"""
+@author:XuMing(xuming624@qq.com)
+@description:
+part of the code from https://github.com/phidatahq/phidata
+"""
+from typing import Optional, Dict, List, Tuple, Any, Union
+
+from agentica.emb.base import Emb
+from agentica.utils.log import logger
+
+try:
+    from huggingface_hub import InferenceClient, SentenceSimilarityInput
+except ImportError:
+    raise ImportError("`huggingface-hub` not installed, please run `pip install huggingface-hub`")
+
+
+class HuggingfaceCustomEmb(Emb):
+    """Huggingface Custom Embedder"""
+
+    model: str = "jinaai/jina-embeddings-v2-base-code"
+    api_key: Optional[str] = None
+    client_params: Optional[Dict[str, Any]] = None
+    huggingface_client: Optional[InferenceClient] = None
+
+    @property
+    def client(self) -> InferenceClient:
+        if self.huggingface_client:
+            return self.huggingface_client
+        _client_params: Dict[str, Any] = {}
+        if self.api_key:
+            _client_params["api_key"] = self.api_key
+        if self.client_params:
+            _client_params.update(self.client_params)
+        self.huggingface_client = InferenceClient(**_client_params)
+        return self.huggingface_client
+
+    def _response(self, text: str):
+        _request_params: SentenceSimilarityInput = {
+            "json": {"inputs": text},
+            "model": self.model,
+        }
+        return self.client.post(**_request_params)
+
+    def get_embedding(self, text: str) -> List[float]:
+        response = self._response(text=text)
+        try:
+            return response
+        except Exception as e:
+            logger.warning(e)
+            return []
+
+    def get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict]]:
+        return super().get_embedding_and_usage(text)

agentica-0.1.4/agentica/emb/ollama_emb.py → agentica-0.1.5/agentica/emb/ollama.py

@@ -4,7 +4,7 @@
 @description:
 part of the code from https://github.com/phidatahq/phidata
 """
-from typing import Optional, Dict, List, Any
+from typing import Optional, Dict, List, Any,Tuple
 
 from agentica.emb.base import Emb
 from agentica.utils.log import logger
@@ -58,3 +58,4 @@ class OllamaEmb(Emb):
 
     def get_embeddings(self, texts: List[str]) -> List[List[float]]:
         return [self.get_embedding(text) for text in texts]
+

agentica-0.1.4/agentica/emb/openai_emb.py → agentica-0.1.5/agentica/emb/openai.py

@@ -5,11 +5,10 @@
 part of the code from https://github.com/phidatahq/phidata
 """
 from os import getenv
-from typing import Optional, Dict, List, Tuple, Any
+from typing import Optional, Dict, List, Tuple, Any, Literal
 
 from openai import OpenAI as OpenAIClient
 from openai.types.create_embedding_response import CreateEmbeddingResponse
-from typing_extensions import Literal
 
 from agentica.emb.base import Emb
 from agentica.utils.log import logger
@@ -41,7 +40,8 @@ class OpenAIEmb(Emb):
             _client_params["base_url"] = self.base_url
         if self.client_params:
             _client_params.update(self.client_params)
-        return OpenAIClient(**_client_params)
+        self.openai_client = OpenAIClient(**_client_params)
+        return self.openai_client
 
     def _response(self, text: str) -> CreateEmbeddingResponse:
         _request_params: Dict[str, Any] = {
@@ -70,7 +70,9 @@ class OpenAIEmb(Emb):
 
         embedding = response.data[0].embedding
         usage = response.usage
-        return embedding, usage.model_dump()
+        if usage:
+            return embedding, usage.model_dump()
+        return embedding, None
 
     def get_embeddings(self, texts: List[str]) -> List[List[float]]:
-        return [self.get_embedding(text) for text in texts]
+        return [self.get_embedding(text) for text in texts]