ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/conf.py

@@ -0,0 +1,395 @@
+import sys
+import os
+from pathlib import Path
+from navconfig import config, BASE_DIR
+from navconfig.logging import logging
+from navigator.conf import default_dsn, CACHE_HOST, CACHE_PORT
+
+
+# # disable debug on some libraries:
+# logging.getLogger(name='httpcore').setLevel(logging.INFO)
+# logging.getLogger(name='httpx').setLevel(logging.INFO)
+# logging.getLogger(name='groq').setLevel(logging.INFO)
+# logging.getLogger(name='selenium.webdriver').setLevel(logging.WARNING)
+# logging.getLogger(name='selenium').setLevel(logging.INFO)
+# logging.getLogger(name='matplotlib').setLevel(logging.WARNING)
+# logging.getLogger(name='PIL').setLevel(logging.INFO)
+logging.getLogger("grpc").setLevel(logging.ERROR)
+os.environ['GRPC_VERBOSITY'] = 'ERROR'
+# logging.getLogger("weasyprint").setLevel(logging.ERROR)  # Suppress WeasyPrint warnings
+# # Suppress tiktoken warnings
+# logging.getLogger("tiktoken").setLevel(logging.ERROR)
+# logging.getLogger("fontTools").setLevel(logging.ERROR)
+
+# Project Root:
+PROJECT_ROOT = BASE_DIR
+# Plugins Directory:
+PLUGINS_DIR = config.get('PLUGINS_DIR', fallback=BASE_DIR.joinpath('plugins'))
+if isinstance(PLUGINS_DIR, str):
+    PLUGINS_DIR = Path(PLUGINS_DIR).resolve()
+if not PLUGINS_DIR.exists():
+    PLUGINS_DIR.mkdir(parents=True, exist_ok=True)
+
+# Static directory
+STATIC_DIR = config.get('STATIC_DIR', fallback=BASE_DIR.joinpath('static'))
+if isinstance(STATIC_DIR, str):
+    STATIC_DIR = Path(STATIC_DIR)
+
+
+# Environment
+ENVIRONMENT = config.get("ENVIRONMENT", fallback="development")
+ENABLE_SWAGGER = config.getboolean("ENABLE_SWAGGER", fallback=True)
+
+
+# Agents Directory
+AGENTS_DIR = config.get('AGENTS_DIR', fallback=BASE_DIR.joinpath('agents'))
+if isinstance(AGENTS_DIR, str):
+    AGENTS_DIR = Path(AGENTS_DIR).resolve()
+if not AGENTS_DIR.exists():
+    AGENTS_DIR.mkdir(parents=True, exist_ok=True)
+
+# Add AGENTS_DIR to sys.path for direct imports (e.g., from agents.troc import ...)
+# Remove if already present to avoid duplicates, then insert at position 0
+# This ensures AGENTS_DIR takes precedence over PLUGINS_DIR even if plugins/__init__.py
+# has already inserted PLUGINS_DIR at position 0
+agents_dir_str = str(AGENTS_DIR)
+if agents_dir_str in sys.path:
+    sys.path.remove(agents_dir_str)
+sys.path.insert(0, agents_dir_str)
+
+
+# MCP Server Directory:
+MCP_SERVER_DIR = config.get(
+    'MCP_SERVER_DIR',
+    fallback=BASE_DIR.joinpath('mcp_servers')
+)
+if isinstance(MCP_SERVER_DIR, str):
+    MCP_SERVER_DIR = Path(MCP_SERVER_DIR).resolve()
+if not MCP_SERVER_DIR.exists():
+    MCP_SERVER_DIR.mkdir(parents=True, exist_ok=True)
+
+# MCP Server defaults
+MCP_SERVER_TRANSPORT = config.get('MCP_SERVER_TRANSPORT', fallback='http')
+MCP_SERVER_HOST = config.get('MCP_SERVER_HOST', fallback='127.0.0.1')
+MCP_SERVER_PORT = config.getint('MCP_SERVER_PORT', fallback=9090)
+MCP_SERVER_NAME = config.get('MCP_SERVER_NAME', fallback='ai-parrot-tools')
+MCP_SERVER_DESCRIPTION = config.get(
+    'MCP_SERVER_DESCRIPTION',
+    fallback='AI-Parrot MCP Tooling'
+)
+MCP_SERVER_LOG_LEVEL = config.get('MCP_SERVER_LOG_LEVEL', fallback='INFO')
+
+# Default tools that should be started with the MCP server
+MCP_STARTED_TOOLS = {
+    # 'MSTeamsToolkit': 'parrot.tools.msteams',
+    # 'PDFPrintTool': 'parrot.tools.pdfprint',
+    'JiraToolkit': 'parrot.tools.jiratoolkit',
+}
+
+# Agents-Bots Prompt directory:
+AGENTS_BOTS_PROMPT_DIR = config.get(
+    'AGENTS_BOTS_PROMPT_DIR',
+    fallback=AGENTS_DIR.joinpath('prompts')
+)
+if isinstance(AGENTS_BOTS_PROMPT_DIR, str):
+    AGENTS_BOTS_PROMPT_DIR = Path(AGENTS_BOTS_PROMPT_DIR).resolve()
+if not AGENTS_BOTS_PROMPT_DIR.exists():
+    AGENTS_BOTS_PROMPT_DIR.mkdir(parents=True, exist_ok=True)
+
+# LLM Model
+DEFAULT_LLM_MODEL_NAME = config.get('LLM_MODEL_NAME', fallback='gemini-2.5-pro')
+
+
+## MILVUS DB ##:
+MILVUS_HOST = config.get('MILVUS_HOST', fallback='localhost')
+MILVUS_PROTOCOL = config.get('MILVUS_PROTOCOL', fallback='http')
+MILVUS_PORT = config.get('MILVUS_PORT', fallback=19530)
+MILVUS_URL = config.get('MILVUS_URL')
+MILVUS_TOKEN = config.get('MILVUS_TOKEN')
+MILVUS_USER = config.get('MILVUS_USER')
+MILVUS_PASSWORD = config.get('MILVUS_PASSWORD')
+MILVUS_SECURE = config.getboolean('MILVUS_SECURE', fallback=False)
+MILVUS_SERVER_NAME = config.get(
+    'MILVUS_SERVER_NAME'
+)
+MILVUS_CA_CERT = config.get('MILVUS_CA_CERT', fallback=None)
+MILVUS_SERVER_CERT = config.get('MILVUS_SERVER_CERT', fallback=None)
+MILVUS_SERVER_KEY = config.get('MILVUS_SERVER_KEY', fallback=None)
+MILVUS_USE_TLSv2 = config.getboolean('MILVUS_USE_TLSv2', fallback=False)
+
+# Postgres Database:
+DBHOST = config.get("DBHOST", fallback="localhost")
+DBUSER = config.get("DBUSER")
+DBPWD = config.get("DBPWD")
+DBNAME = config.get("DBNAME", fallback="navigator")
+DBPORT = config.get("DBPORT", fallback=5432)
+# sqlalchemy+asyncpg connector:
+default_sqlalchemy_pg = f"postgresql+asyncpg://{DBUSER}:{DBPWD}@{DBHOST}:{DBPORT}/{DBNAME}"
+
+PG_USER = config.get('PG_USER', fallback=DBUSER)
+PG_PWD = config.get('PG_PWD', fallback=DBPWD)
+PG_HOST = config.get('PG_HOST', fallback=DBHOST)
+PG_PORT = config.get('PG_PORT', fallback=DBPORT)
+PG_DATABASE = config.get('PG_DATABASE', fallback=DBNAME)
+# asyncpg url for sqlalchemy:
+asyncpg_sqlalchemy_url = f"postgresql+asyncpg://{PG_USER}:{PG_PWD}@{PG_HOST}:{PG_PORT}/{PG_DATABASE}"
+
+# ScyllaDB Database:
+SCYLLADB_DRIVER = config.get('SCYLLADB_DRIVER', fallback='scylladb')
+SCYLLADB_HOST = config.get('SCYLLADB_HOST', fallback='localhost')
+SCYLLADB_PORT = config.getint('SCYLLADB_PORT', fallback=9042)
+SCYLLADB_USERNAME = config.get('SCYLLADB_USERNAME', fallback='navigator')
+SCYLLADB_PASSWORD = config.get('SCYLLADB_PASSWORD', fallback='navigator')
+SCYLLADB_KEYSPACE = config.get('SCYLLADB_KEYSPACE', fallback='navigator')
+
+
+# BigQuery Configuration:
+BIGQUERY_CREDENTIALS = config.get('BIGQUERY_CREDENTIALS')
+BIGQUERY_PROJECT_ID = config.get('BIGQUERY_PROJECT_ID', fallback='navigator')
+BIGQUERY_DATASET = config.get('BIGQUERY_DATASET', fallback='navigator')
+
+# Redis History Configuration:
+REDIS_HOST = config.get('REDIS_HOST', fallback='localhost')
+REDIS_PORT = config.get('REDIS_PORT', fallback=6379)
+REDIS_DB = config.get('REDIS_DB', fallback=1)
+REDIS_URL = config.get('REDIS_URL', fallback=f"redis://{REDIS_HOST}:{REDIS_PORT}/{REDIS_DB}")
+REDIS_HISTORY_DB = config.get('REDIS_HISTORY_DB', fallback=3)
+REDIS_HISTORY_URL = f"redis://{REDIS_HOST}:{REDIS_PORT}/{REDIS_HISTORY_DB}"
+REDIS_SERVICES_URL = f"redis://{REDIS_HOST}:{REDIS_PORT}/4"
+
+def resolve_cert(crt):
+    cert = Path(crt)
+    if not cert.is_absolute():
+        cert = BASE_DIR.joinpath(cert)
+    else:
+        cert.resolve()
+    return cert
+
+if MILVUS_SERVER_CERT:
+    MILVUS_SERVER_CERT = str(resolve_cert(MILVUS_SERVER_CERT))
+if MILVUS_CA_CERT:
+    MILVUS_CA_CERT = str(resolve_cert(MILVUS_CA_CERT))
+if MILVUS_SERVER_KEY:
+    MILVUS_SERVER_KEY = str(resolve_cert(MILVUS_SERVER_KEY))
+
+# QDRANT:
+QDRANT_PROTOCOL = config.get('QDRANT_PROTOCOL', fallback='http')
+QDRANT_HOST = config.get('QDRANT_HOST', fallback='localhost')
+QDRANT_PORT = config.get('QDRANT_PORT', fallback=6333)
+QDRANT_USE_HTTPS = config.getboolean('QDRANT_USE_HTTPS', fallback=False)
+QDRANT_URL = config.get('QDRANT_URL')
+# QDRANT Connection Type: server or cloud
+QDRANT_CONN_TYPE = config.get('QDRANT_CONN_TYPE', fallback='server')
+
+# ChromaDB:
+CHROMADB_HOST = config.get('CHROMADB_HOST', fallback='localhost')
+CHROMADB_PORT = config.get('CHROMADB_PORT', fallback=8000)
+
+# Embedding Device:
+EMBEDDING_DEVICE = config.get('EMBEDDING_DEVICE', fallback='cpu')
+EMBEDDING_DEFAULT_MODEL = config.get(
+    'EMBEDDING_DEFAULT_MODEL',
+    fallback='sentence-transformers/all-MiniLM-L12-v2'
+)
+KB_DEFAULT_MODEL = config.get(
+    'KB_DEFAULT_MODEL',
+    fallback='sentence-transformers/paraphrase-MiniLM-L3-v2'
+)
+HUGGINGFACE_EMBEDDING_CACHE_DIR = config.get(
+    'HUGGINGFACE_EMBEDDING_CACHE_DIR',
+    fallback=BASE_DIR.joinpath('model_cache', 'huggingface')
+)
+HUGGINGFACEHUB_API_TOKEN = config.get('HUGGINGFACEHUB_API_TOKEN')
+MAX_VRAM_AVAILABLE = config.get('MAX_VRAM_AVAILABLE', fallback=20000)
+RAM_AVAILABLE = config.get('RAM_AVAILABLE', fallback=819200)
+CUDA_DEFAULT_DEVICE = config.get('CUDA_DEFAULT_DEVICE', fallback='cpu')
+CUDA_DEFAULT_DEVICE_NUMBER = config.getint('CUDA_DEFAULT_DEVICE_NUMBER', fallback=0)
+MAX_BATCH_SIZE = config.get('MAX_BATCH_SIZE', fallback=2048)
+
+# Enable Teams Bot:
+ENABLE_AZURE_BOT = config.getboolean('ENABLE_AZURE_BOT', fallback=True)
+
+## Google Services:
+GOOGLE_API_KEY = config.get('GOOGLE_API_KEY')
+### Google Service Credentials:
+GA_SERVICE_ACCOUNT_NAME = config.get('GA_SERVICE_ACCOUNT_NAME', fallback="google.json")
+GA_SERVICE_PATH = config.get('GA_SERVICE_PATH', fallback="env/google/")
+if isinstance(GA_SERVICE_PATH, str):
+    GA_SERVICE_PATH = Path(GA_SERVICE_PATH)
+
+GOOGLE_TTS_SERVICE = config.get(
+    'GOOGLE_TTS_SERVICE',
+    fallback=GA_SERVICE_PATH.joinpath('tts-service.json')
+)
+if isinstance(GOOGLE_TTS_SERVICE, str):
+    GOOGLE_TTS_SERVICE = Path(GOOGLE_TTS_SERVICE)
+if not GOOGLE_TTS_SERVICE.is_absolute():
+    GOOGLE_TTS_SERVICE = BASE_DIR.joinpath(GOOGLE_TTS_SERVICE)
+if not GOOGLE_TTS_SERVICE.exists():
+    GOOGLE_TTS_SERVICE = None
+
+# BASE STATIC:
+BASE_STATIC_URL = config.get(
+    'BASE_STATIC_URL',
+    fallback='http://localhost:5000/static'
+)
+
+# Google SerpAPI:
+SERPAPI_API_KEY = config.get('SERPAPI_API_KEY')
+
+# Groq API Key:
+GROQ_API_KEY = config.get('GROQ_API_KEY')
+DEFAULT_GROQ_MODEL = config.get('DEFAULT_GROQ_MODEL', fallback='qwen/qwen3-32b')
+
+# Ethical Principle:
+ETHICAL_PRINCIPLE = config.get(
+    'ETHICAL_PRINCIPLE',
+    fallback='The model should only talk about ethical and legal things.'
+)
+
+# Embedding Configuration:
+
+# VERTEX
+VERTEX_PROJECT_ID = config.get('VERTEX_PROJECT_ID')
+VERTEX_REGION = config.get('VERTEX_REGION')
+
+# OpenAI:
+OPENAI_API_KEY = config.get('OPENAI_API_KEY')
+OPENAI_ORGANIZATION = config.get('OPENAI_ORGANIZATION')
+
+## HTTPClioent
+HTTPCLIENT_MAX_SEMAPHORE = config.getint("HTTPCLIENT_MAX_SEMAPHORE", fallback=5)
+HTTPCLIENT_MAX_WORKERS = config.getint("HTTPCLIENT_MAX_WORKERS", fallback=1)
+
+## Google API:
+GOOGLE_API_KEY = config.get('GOOGLE_API_KEY')
+GOOGLE_SEARCH_API_KEY = config.get('GOOGLE_SEARCH_API_KEY')
+GOOGLE_SEARCH_ENGINE_ID = config.get('GOOGLE_SEARCH_ENGINE_ID')
+GOOGLE_PLACES_API_KEY = config.get('GOOGLE_PLACES_API_KEY')
+GOOGLE_CREDENTIALS_FILE = Path(
+    config.get(
+        'GOOGLE_CREDENTIALS_FILE',
+        fallback=BASE_DIR.joinpath('env', 'google', 'key.json')
+    )
+)
+
+## LLM default config:
+DEFAULT_LLM_MODEL = config.get('LLM_MODEL', fallback='gemini-2.5-flash')
+DEFAULT_LLM_TEMPERATURE = config.get('LLM_TEMPERATURE', fallback=0.1)
+
+"""
+Amazon AWS Credentials
+"""
+aws_region = config.get("AWS_REGION", fallback="us-east-1")
+aws_bucket = config.get("AWS_BUCKET", fallback="static-files")
+aws_key = config.get("AWS_KEY")
+aws_secret = config.get("AWS_SECRET")
+
+AWS_ACCESS_KEY = config.get("AWS_ACCESS_KEY", fallback=aws_key)
+AWS_SECRET_KEY = config.get("AWS_SECRET_KEY", fallback=aws_secret)
+AWS_REGION_NAME = config.get("AWS_REGION_NAME", fallback=aws_region)
+AWS_DEFAULT_CLOUDWATCH_LOG_GROUP = config.get("AWS_DEFAULT_CLOUDWATCH_LOG_GROUP", fallback="/parrot/logs")
+
+AWS_CREDENTIALS = {
+    "default": {
+        "use_credentials": config.get("aws_credentials", fallback=False),
+        "aws_key": aws_key,
+        "aws_secret": aws_secret,
+        "region_name": aws_region,
+        "bucket_name": aws_bucket,
+    },
+    "monitoring": {
+        "use_credentials": config.get("aws_monitor_credentials", fallback=True),
+        "aws_key": AWS_ACCESS_KEY,
+        "aws_secret": AWS_SECRET_KEY,
+        "region_name": AWS_REGION_NAME,
+    },
+    "cloudwatch": {
+        "use_credentials": True,
+        "aws_key": config.get("AWS_CLOUDWATCH_KEY"),
+        "aws_secret": config.get("AWS_CLOUDWATCH_SECRET"),
+        "region_name": config.get("AWS_CLOUDWATCH_REGION", fallback="us-east-1"),
+    }
+}
+
+## Tools:
+OPENWEATHER_APPID = config.get('OPENWEATHER_APPID')
+
+# NOTIFICATIONS:
+TEAMS_NOTIFY_TENANT_ID = config.get("TEAMS_NOTIFY_TENANT_ID")
+TEAMS_NOTIFY_CLIENT_ID = config.get("TEAMS_NOTIFY_CLIENT_ID")
+TEAMS_NOTIFY_CLIENT_SECRET = config.get("TEAMS_NOTIFY_CLIENT_SECRET")
+TEAMS_NOTIFY_USERNAME = config.get("TEAMS_NOTIFY_USERNAME")
+TEAMS_NOTIFY_PASSWORD = config.get("TEAMS_NOTIFY_PASSWORD")
+MS_TEAMS_DEFAULT_TEAMS_ID = config.get("MS_TEAMS_DEFAULT_TEAMS_ID")
+MS_TEAMS_DEFAULT_CHANNEL_ID = config.get("MS_TEAMS_DEFAULT_CHANNEL_ID")
+
+## MS Teams Toolkit:
+MS_TEAMS_CLIENT_SECRET = config.get('MS_TEAMS_CLIENT_SECRET')
+MS_TEAMS_CLIENT_ID = config.get('MS_TEAMS_CLIENT_ID')
+MS_TEAMS_TENANT_ID = config.get('MS_TEAMS_TENANT_ID')
+MS_TEAMS_USERNAME = config.get('TEAMS_NOTIFY_USERNAME')
+MS_TEAMS_PASSWORD = config.get('TEAMS_NOTIFY_PASSWORD')
+
+## Office 365:
+O365_CLIENT_ID = config.get('O365_CLIENT_ID')
+O365_CLIENT_SECRET = config.get('O365_CLIENT_SECRET')
+O365_TENANT_ID = config.get('O365_TENANT_ID')
+
+# Sharepoint:
+SHAREPOINT_APP_ID = config.get('SHAREPOINT_APP_ID')
+SHAREPOINT_APP_SECRET = config.get('SHAREPOINT_APP_SECRET')
+SHAREPOINT_TENANT_ID = config.get('SHAREPOINT_TENANT_ID')
+SHAREPOINT_TENANT_NAME = config.get('SHAREPOINT_TENANT_NAME')
+SHAREPOINT_SITE_ID = config.get('SHAREPOINT_SITE_ID')
+SHAREPOINT_DEFAULT_HOST = config.get('SHAREPOINT_DEFAULT_HOST')
+
+# Employee Hierarchy Configuration:
+EMPLOYEES_TABLE = config.get('EMPLOYEES_TABLE', fallback='troc.troc_employees')
+
+# Workday SOAP settings
+WORKDAY_DEFAULT_TENANT = config.get('WORKDAY_DEFAULT_TENANT', fallback='nav')
+WORKDAY_CLIENT_ID = config.get("WORKDAY_CLIENT_ID")
+WORKDAY_CLIENT_SECRET = config.get("WORKDAY_CLIENT_SECRET")
+WORKDAY_TOKEN_URL = config.get("WORKDAY_TOKEN_URL")
+WORKDAY_WSDL_PATH = config.get(
+    "WORKDAY_WSDL_PATH",
+    fallback=BASE_DIR.joinpath("env", "workday", "staffing_custom_44_2.wsdl")
+)
+WORKDAY_WSDL_TIME = config.get(
+    "WORKDAY_WSDL_TIME",
+    fallback=BASE_DIR.joinpath("env", "workday", "timetracking_custom_44_2.wsdl")
+)
+WORKDAY_WSDL_HUMAN_RESOURCES = config.get(
+    "WORKDAY_WSDL_HUMAN_RESOURCES",
+    fallback=BASE_DIR.joinpath("env", "workday", "humanresources_troc_44_2.wsdl")
+)
+WORKDAY_WSDL_FINANCIAL_MANAGEMENT = config.get(
+    "WORKDAY_WSDL_FINANCIAL_MANAGEMENT",
+    fallback=BASE_DIR.joinpath("env", "workday", "financial_management_45.wsdl")
+)
+WORKDAY_WSDL_RECRUITING = config.get(
+    "WORKDAY_WSDL_RECRUITING",
+    fallback=BASE_DIR.joinpath("env", "workday", "recruiting_44_2.wsdl")
+)
+WORKDAY_WSDL_ABSENCE_MANAGEMENT = config.get(
+    "WORKDAY_WSDL_ABSENCE_MANAGEMENT",
+    fallback=BASE_DIR.joinpath("env", "workday", "absence_management_45_custom.wsdl")
+)
+WORKDAY_REFRESH_TOKEN = config.get("WORKDAY_REFRESH_TOKEN", fallback=None)
+
+WORKDAY_WSDL_PATHS = {
+    "human_resources": WORKDAY_WSDL_HUMAN_RESOURCES,
+    "absence_management": WORKDAY_WSDL_ABSENCE_MANAGEMENT,
+    "time_tracking": WORKDAY_WSDL_TIME,
+    "staffing": WORKDAY_WSDL_PATH,
+    "financial_management": WORKDAY_WSDL_FINANCIAL_MANAGEMENT,
+    "recruiting": WORKDAY_WSDL_RECRUITING
+}
+
+# Final sys.path adjustment: Ensure AGENTS_DIR takes precedence over PLUGINS_DIR
+# This is necessary because parrot.plugins.__init__.py may have inserted PLUGINS_DIR
+# at position 0 during module loading (after our initial AGENTS_DIR insertion above)
+if agents_dir_str in sys.path:
+    sys.path.remove(agents_dir_str)
+sys.path.insert(0, agents_dir_str)

parrot/embeddings/__init__.py

@@ -0,0 +1,9 @@
+from .huggingface import SentenceTransformerModel
+from .google import GoogleEmbeddingModel
+from .openai import OpenAIEmbeddingModel
+
+supported_embeddings = {
+    'huggingface': 'SentenceTransformerModel',
+    'google': 'GoogleEmbeddingModel',
+    'openai': 'OpenAIEmbeddingModel',
+}

parrot/embeddings/base.py

@@ -0,0 +1,157 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional, Union, Any
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+import numpy as np
+from navconfig.logging import logging
+from ..conf import (
+    EMBEDDING_DEVICE,
+    CUDA_DEFAULT_DEVICE
+)
+
+
+class EmbeddingModel(ABC):
+    """
+    Abstract base class for embedding models.
+    It ensures that embedding models can be used interchangeably.
+    """
+    def __init__(self, model_name: str, **kwargs):
+        self.model_name = model_name
+        self.logger = logging.getLogger(f"parrot.{self.__class__.__name__}")
+        self.executor = ThreadPoolExecutor(max_workers=4)
+        self._model_lock = asyncio.Lock()
+        self._dimension = None
+        # Lazy initialization
+        self._model = None
+        self._device = None
+        self._kwargs = kwargs
+
+    @property
+    def device(self):
+        if self._device is None:
+            _, self._device, self._dtype = self._get_device()
+        return self._device
+
+    @property
+    def model(self):
+        if self._model is None:
+            self._model = self._create_embedding(
+                model_name=self.model_name,
+                **self._kwargs
+            )
+        return self._model
+
+    def _get_device(
+        self,
+        device_type: str = None,
+        cuda_number: int = 0
+    ):
+        """Get Default device for Torch and transformers.
+
+        """
+        import torch
+        dev = torch.device("cpu")
+        pipe_dev = -1
+        dtype = torch.float32
+        if device_type == 'cpu':
+            pipe_dev = torch.device('cpu')
+        if CUDA_DEFAULT_DEVICE == 'cpu':
+            # Use CPU forced
+            pipe_dev = torch.device('cpu')
+        if torch.cuda.is_available():
+            dev = torch.device("cuda")
+            pipe_dev = 0  # first GPU
+            # prefer bf16 if supported; else fp16
+            if torch.cuda.is_bf16_supported():
+                dtype = torch.bfloat16
+            else:
+                dtype = torch.float16
+            pipe_dev = torch.device(f'cuda:{cuda_number}')
+        if device_type == 'cuda':
+            if torch.cuda.is_bf16_supported():
+                dtype = torch.bfloat16
+            else:
+                dtype = torch.float16
+            pipe_dev = torch.device(f'cuda:{cuda_number}')
+        if torch.backends.mps.is_available():
+            # Use CUDA Multi-Processing Service if available
+            dev = torch.device("mps")
+            pipe_dev = torch.device("mps")
+            dtype = torch.float32  # fp16 on MPS is still flaky
+        else:
+            pipe_dev = torch.device(EMBEDDING_DEVICE)
+        return pipe_dev, dev, dtype
+
+    def get_embedding_dimension(self) -> int:
+        return self._dimension
+
+    async def initialize_model(self):
+        """Async model initialization with GPU optimization"""
+        async with self._model_lock:
+            if self.model is None:
+                loop = asyncio.get_event_loop()
+                self.model = await loop.run_in_executor(
+                    self.executor,
+                    self._create_embedding
+                )
+
+
+    @abstractmethod
+    def _create_embedding(self, model_name: str, **kwargs) -> Any:
+        """
+        Loads and returns the embedding model instance.
+        """
+        pass
+
+    def embed_documents(
+        self,
+        texts: List[str],
+        batch_size: Optional[int] = None
+    ) -> List[List[float]]:
+        """
+        Generates embeddings for a list of documents.
+
+        Args:
+            texts: A list of document strings.
+
+        Returns:
+            A list of embedding vectors.
+        """
+        return self.model.encode(texts, convert_to_tensor=False).tolist()
+
+    def embed_query(
+        self,
+        text: str,
+        as_nparray: bool = False
+    ) -> Union[List[float], List[np.ndarray]]:
+        """
+        Generates an embedding for a single query string.
+
+        Args:
+            text: The query string.
+
+        Returns:
+            The embedding vector for the query.
+        """
+        embeddings = self.model.encode(
+            text,
+            convert_to_tensor=False,
+            normalize_embeddings=True,
+            show_progress_bar=False
+        )
+        if as_nparray:
+            return np.vstack(embeddings)
+        return embeddings.tolist()
+
+    def free(self):
+        """
+        Frees up resources used by the model.
+        """
+        import torch
+        self.model = None
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    @abstractmethod
+    async def encode(self, texts: List[str], **kwargs) -> np.ndarray:
+        pass

parrot/embeddings/google.py

@@ -0,0 +1,98 @@
+from typing import Any, List, Union
+import numpy as np
+from navconfig import config
+# This file is part of Parrot, an open-source project.
+from .base import EmbeddingModel
+
+class GoogleEmbeddingModel(EmbeddingModel):
+    """A wrapper class for Google Embedding models using the Gemini API.
+    """
+    model_name: str = "gemini-embedding-001"
+
+    def __init__(self, model_name: str = None, output_dimensionality: int = None, **kwargs):
+        self.api_key = kwargs.pop('api_key', config.get('GOOGLE_API_KEY'))
+        if model_name:
+            self.model_name = model_name
+        self.output_dimensionality = output_dimensionality
+        super().__init__(**kwargs)
+
+    def _create_embedding(self, model_name: str = None, **kwargs) -> Any:
+        """
+        Creates and returns a Google Embedding model instance.
+
+        Args:
+            model_name: The name of the Google model to load.
+
+        Returns:
+            An instance of Google Embedding model.
+        """
+        from google import genai
+        if model_name:
+            self.model_name = model_name
+        self.logger.info(
+            f"Loading embedding model '{self.model_name}'"
+        )
+        self.client = genai.Client(api_key=self.api_key)
+        return self.client
+
+    def _normalize_embeddings(self, embeddings: List[Any]) -> List[List[float]]:
+        if not self.output_dimensionality or self.output_dimensionality == 3072:
+            return [e.values for e in embeddings]
+        
+        # Normalize embeddings for lower dimensions
+        normalized = []
+        for embedding in embeddings:
+            if hasattr(embedding, 'values'):
+                val = np.array(embedding.values)
+            else:
+                val = np.array(embedding)
+            norm = np.linalg.norm(val)
+            if norm > 0:
+                val = val / norm
+            normalized.append(val.tolist())
+        return normalized
+
+    async def encode(self, texts: List[str], **kwargs) -> List[List[float]]:
+        from google.genai import types
+        
+        call_kwargs = {
+            "model": self.model_name,
+            "contents": texts
+        }
+        if self.output_dimensionality:
+            call_kwargs["config"] = types.EmbedContentConfig(
+                output_dimensionality=self.output_dimensionality
+            )
+            
+        result = self.client.models.embed_content(**call_kwargs)
+        if self.output_dimensionality:
+            return self._normalize_embeddings(result.embeddings)
+        return [e.values for e in result.embeddings]
+
+    def embed_query(
+        self,
+        text: str,
+        as_nparray: bool = False
+    ) -> Union[List[float], List[np.ndarray]]:
+        from google.genai import types
+        
+        call_kwargs = {
+            "model": self.model_name,
+            "contents": [text]
+        }
+        if self.output_dimensionality:
+            call_kwargs["config"] = types.EmbedContentConfig(
+                output_dimensionality=self.output_dimensionality
+            )
+
+        result = self.client.models.embed_content(**call_kwargs)
+        
+        embeddings = result.embeddings
+        if self.output_dimensionality:
+             embeddings = self._normalize_embeddings(embeddings)
+        else:
+             embeddings = [e.values for e in embeddings]
+
+        if as_nparray:
+            return [np.array(embedding) for embedding in embeddings]
+        return embeddings