Merge remote-tracking branch 'upstream/main'

Author: chiruu12

.github/workflows/main.yaml

@@ -40,6 +40,7 @@ jobs:
           # - "3.13"
         os:
           - "ubuntu"
+          - "macos"
 
     runs-on: ${{ matrix.os }}-latest
     timeout-minutes: 20

conda/dev.yaml

@@ -3,7 +3,7 @@ channels:
   - nodefaults
   - conda-forge
 dependencies:
-  - python <3.13
+  - python 3.9.*
   - pip
   - poetry >=2
   - nodejs # used by semantic-release

docs/changelog.md

@@ -1,6 +1,19 @@
 # Release Notes
 ---
 
+# [0.13.0](https://github.com/osl-incubator/rago/compare/0.12.0...0.13.0) (2025-03-13)
+
+
+### Bug Fixes
+
+* **pkg:** Add support for MacOS ([#75](https://github.com/osl-incubator/rago/issues/75)) ([e6a33b0](https://github.com/osl-incubator/rago/commit/e6a33b0d967d21130cf005789989dc9e1c28c4fb))
+
+
+### Features
+
+* Add cohere backend support ([#62](https://github.com/osl-incubator/rago/issues/62)) ([6817ba0](https://github.com/osl-incubator/rago/commit/6817ba08607e5366b4b36dba7b91644eced7edb7))
+* **generation:** add backend for DeepSeek's generation class ([#49](https://github.com/osl-incubator/rago/issues/49)) ([47947d6](https://github.com/osl-incubator/rago/commit/47947d65105db88c7a021ae7da8b48cff8ce58d1))
+
 # [0.12.0](https://github.com/osl-incubator/rago/compare/0.11.3...0.12.0) (2025-02-11)
 
 

poetry.lock

@@ -1,6 +1,6 @@
 [project]
 name = "rago"
-version = "0.12.0"  # semantic-release
+version = "0.13.0"  # semantic-release
 description = "Rago is a lightweight framework for RAG"
 readme = "README.md"
 authors = [
@@ -19,6 +19,8 @@ requires-python = ">=3.9,<3.13"
 
 dependencies = [
   "typeguard >=4.0",
+  "numpy >=2,<=2.2 ; python_version > '3.9'",
+  "numpy >=2,<2.1 ; python_version < '3.10'",
   "faiss-cpu >=1.9.0",
   "sentencepiece >=0.2.0",
   "sentence-transformers >=3.2.0",
@@ -29,39 +31,37 @@ dependencies = [
   "pypdf >=5",
   "langchain >=0.3.7",
   "langchain-community >=0.3.7",
-  "spacy >=3",
+  "spacy >=3.8.0 ; python_version > '3.9'",
+  "spacy >=3.8.0,<3.8.4 ; python_version < '3.10'",
   "instructor >=1",
   "pydantic >=2",
-  "torch >=2.5",
-  "torchvision >=0.20",
+  "torch >=2,<2.6",
   "eval-type-backport >=0.2 ; python_version < '3.10'",
   "joblib (>=1.4.2,<2.0.0)",
-  "cohere >=5.13.4"
+  "cohere >=5.13.4",
+  "fireworks-ai>0.15.10"
 ]
 
 [build-system]
 requires = ["poetry-core>=2", "poetry>=2"]
 build-backend = "poetry.core.masonry.api"
 
-[tool.poetry.extras]
-cpu = ["torch", "torchvision"]
-gpu = ["torch", "torchvision"]
+# [project.optional-dependencies]
+# cpu = ["torch"]
+# gpu = ["torch"]
 
-[[tool.poetry.source]]
-name = "pytorch-cpu"
-url = "https://download.pytorch.org/whl/cpu"
-priority = "supplemental"
+# [[tool.poetry.source]]
+# name = "pytorch-cpu"
+# url = "https://download.pytorch.org/whl/cpu"
+# priority = "explicit"
 
-[tool.poetry.dependencies]
-python = ">=3.9,<3.13"
-torch = [
-  {version = ">=2.5.0", markers="extra!='gpu'", source="pytorch-cpu"},
-  {version = ">=2.5.0", markers="extra=='gpu'"},
-]
-torchvision = [
-  {version = ">=0.20.0", markers="extra!='gpu'", source="pytorch-cpu"},
-  {version = ">=0.20.0", markers="extra=='gpu'"},
-]
+# [tool.poetry.dependencies]
+# python = ">=3.9,<3.13"
+# torch = [
+#   { version = ">=2,<2.6", platform = "darwin" },
+#   { version = ">=2,<2.6", platform = "linux", source="pytorch-cpu", markers="extra!='gpu'" },
+#   { version = ">=2,<2.6", platform = "linux", markers="extra=='gpu'" },
+# ]
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=7.3.2"

pyproject.toml

@@ -10,7 +10,7 @@ def get_version() -> str:
     try:
         return importlib_metadata.version(__name__)
     except importlib_metadata.PackageNotFoundError:  # pragma: no cover
-        return '0.12.0'  # semantic-release
+        return '0.13.0'  # semantic-release
 
 
 version = get_version()

src/rago/__init__.py

@@ -4,13 +4,15 @@
 
 from rago.augmented.base import AugmentedBase
 from rago.augmented.cohere import CohereAug
+from rago.augmented.fireworks import FireworksAug
 from rago.augmented.openai import OpenAIAug
 from rago.augmented.sentence_transformer import SentenceTransformerAug
 from rago.augmented.spacy import SpaCyAug
 
 __all__ = [
     'AugmentedBase',
     'CohereAug',
+    'FireworksAug',
     'OpenAIAug',
     'SentenceTransformerAug',
     'SpaCyAug',

src/rago/augmented/__init__.py

@@ -0,0 +1,75 @@
+"""Classes for augmentation with Fireworks embeddings."""
+
+from __future__ import annotations
+
+from hashlib import sha256
+from typing import cast
+
+import numpy as np
+import openai  # fireworks client doesnt have query
+
+
+# embeddings model feature yet
+from typeguard import typechecked
+
+from rago.augmented.base import AugmentedBase, EmbeddingType
+
+
+@typechecked
+class FireworksAug(AugmentedBase):
+    """Class for augmentation with Fireworks embeddings."""
+
+    default_model_name = 'nomic-ai/nomic-embed-text-v1.5'  # embedding model
+    default_top_k = 3
+
+    def _setup(self) -> None:
+        """Set up the object with initial parameters."""
+        if not self.api_key:
+            raise ValueError('API key for Fireworks is required.')
+        self.openai_client = openai.OpenAI(
+            base_url='https://api.fireworks.ai/inference/v1',
+            api_key=self.api_key,
+        )
+
+    def get_embedding(self, content: list[str]) -> EmbeddingType:
+        """Retrieve the embedding for given texts using the OpenAI client."""
+        cache_key = sha256(''.join(content).encode('utf-8')).hexdigest()
+        cached = self._get_cache(cache_key)
+        if cached is not None:
+            return cast(EmbeddingType, cached)
+
+        # Using the OpenAI embeddings API call for fireworks
+        response = self.openai_client.embeddings.create(
+            model=self.model_name,
+            input=content,
+        )
+        result = np.array(
+            [data.embedding for data in response.data], dtype=np.float32
+        )
+        self._save_cache(cache_key, result)
+        return result
+
+    def search(
+        self, query: str, documents: list[str], top_k: int = 0
+    ) -> list[str]:
+        """Search an encoded query into vector database."""
+        if not hasattr(self, 'db') or not self.db:
+            raise Exception('Vector database (db) is not initialized.')
+
+        document_encoded = self.get_embedding(documents)
+        query_encoded = self.get_embedding([query])
+        top_k = top_k or self.top_k or self.default_top_k or 1
+
+        self.db.embed(document_encoded)
+        scores, indices = self.db.search(query_encoded, top_k=top_k)
+
+        self.logs['indices'] = indices
+        self.logs['scores'] = scores
+        self.logs['search_params'] = {
+            'query_encoded': query_encoded,
+            'top_k': top_k,
+        }
+
+        retrieved_docs = [documents[i] for i in indices if i >= 0]
+
+        return retrieved_docs

src/rago/augmented/fireworks.py

@@ -5,6 +5,7 @@
 from rago.generation.base import GenerationBase
 from rago.generation.cohere import CohereGen
 from rago.generation.deepseek import DeepSeekGen
+from rago.generation.fireworks import FireworksGen
 from rago.generation.gemini import GeminiGen
 from rago.generation.hugging_face import HuggingFaceGen
 from rago.generation.llama import LlamaGen
@@ -13,6 +14,7 @@
 __all__ = [
     'CohereGen',
     'DeepSeekGen',
+    'FireworksGen',
     'GeminiGen',
     'GenerationBase',
     'HuggingFaceGen',

src/rago/generation/__init__.py

@@ -0,0 +1,67 @@
+"""FireworksGen class for text generation using Fireworks API."""
+
+from __future__ import annotations
+
+from typing import cast
+
+import instructor
+
+from fireworks.client import Fireworks
+from pydantic import BaseModel
+from typeguard import typechecked
+
+from rago.generation.base import GenerationBase
+
+
+@typechecked
+class FireworksGen(GenerationBase):
+    """Fireworks AI generation model for text generation."""
+
+    default_model_name: str = 'accounts/fireworks/models/llama-v3-8b-instruct'
+    default_api_params = {  # noqa: RUF012
+        'top_p': 0.9,
+    }
+
+    def _setup(self) -> None:
+        """Set up the object with the initial parameters."""
+        model = Fireworks(api_key=self.api_key)
+
+        self.model = (
+            instructor.from_fireworks(
+                client=model,
+                mode=instructor.Mode.FIREWORKS_JSON,
+            )
+            if self.structured_output
+            else model
+        )
+
+    def generate(self, query: str, context: list[str]) -> str | BaseModel:
+        """Generate text using Fireworks AI's API."""
+        input_text = self.prompt_template.format(
+            query=query, context=' '.join(context)
+        )
+
+        api_params = self.api_params or self.default_api_params
+
+        messages = []
+        if self.system_message:
+            messages.append({'role': 'system', 'content': self.system_message})
+        messages.append({'role': 'user', 'content': input_text})
+
+        model_params = {
+            'model': self.model_name,
+            'messages': messages,
+            'max_tokens': self.output_max_length,
+            'temperature': self.temperature,
+            **api_params,
+        }
+
+        if self.structured_output:
+            model_params['response_model'] = self.structured_output
+            response = self.model.chat.completions.create(**model_params)
+            self.logs['model_params'] = model_params
+            return cast(BaseModel, response)
+
+        response = self.model.chat.completions.create(**model_params)
+        self.logs['model_params'] = model_params
+        return cast(str, response.choices[0].message.content.strip())

src/rago/generation/fireworks.py

@@ -3,5 +3,5 @@ OPENAI_API_KEY=${OPENAI_API_KEY}
 GEMINI_API_KEY=${GEMINI_API_KEY}
 TOKENIZERS_PARALLELISM=false
 COHERE_API_KEY=${COHERE_API_KEY}
-#FIREWORKS_API_KEY=${FIREWORKS_API_KEY}
+FIREWORKS_API_KEY=${FIREWORKS_API_KEY}
 #TOGETHER_API_KEY=${TOGETHER_API_KEY}

tests/.env.tpl

@@ -74,3 +74,14 @@ def api_key_cohere(env) -> str:
             'Please set the COHERE_API_KEY environment variable.'
         )
     return key
+
+
+@pytest.fixture
+def api_key_fireworks(env) -> str:
+    """Fixture for Fireworks API key from environment."""
+    key = os.getenv('FIREWORKS_API_KEY')
+    if not key:
+        raise EnvironmentError(
+            'Please set the FIREWORKS_API_KEY environment variable.'
+        )
+    return key

tests/conftest.py

@@ -4,11 +4,12 @@
 
 import pytest
 
-from rago.augmented import CohereAug, OpenAIAug, SpaCyAug
+from rago.augmented import CohereAug, FireworksAug, OpenAIAug, SpaCyAug
 
 API_MAP = {
     OpenAIAug: 'api_key_openai',
     CohereAug: 'api_key_cohere',
+    FireworksAug: 'api_key_fireworks',
 }
 
 gen_models = [
@@ -26,9 +27,14 @@
             model_name='text-embedding-3-small',
         ),
     ),
+    # model 2
     partial(
         CohereAug,
     ),
+    # model 3
+    partial(
+        FireworksAug,
+    ),
 ]
 
 
@@ -52,6 +58,7 @@ def test_aug_spacy(
     api_key_openai: str,
     api_key_cohere: str,
     api_key_gemini: str,
+    api_key_fireworks: str,
     api_key_hugging_face: str,
     partial_model: partial,
 ) -> None: