feat: Add support for fireworks-ai (osl-incubator#74)

Author: chiruu12

poetry.lock

@@ -38,7 +38,8 @@ dependencies = [
   "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]

pyproject.toml

@@ -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: