fix gemini

Author: xmnlab

src/rago/core.py

@@ -4,6 +4,7 @@
 
 from typing import Any
 
+from pydantic import BaseModel
 from typeguard import typechecked
 
 from rago.augmented.base import AugmentedBase
@@ -50,7 +51,7 @@ def __init__(
             'generation': generation.logs,
         }
 
-    def prompt(self, query: str, device: str = 'auto') -> str:
+    def prompt(self, query: str, device: str = 'auto') -> str | BaseModel:
         """Run the pipeline for a specific prompt.
 
         Parameters
@@ -72,7 +73,7 @@ def prompt(self, query: str, device: str = 'auto') -> str:
         aug_data = self.augmented.search(query, ret_data)
         self.logs['augmented']['result'] = aug_data
 
-        gen_data: str = self.generation.generate(query, context=aug_data)
+        gen_data = self.generation.generate(query, context=aug_data)
         self.logs['generation']['result'] = gen_data
 
         return gen_data

src/rago/generation/base.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import Any, Optional
+from typing import Any, Optional, Type
 
 import torch
 
@@ -27,7 +27,7 @@ class GenerationBase:
     prompt_template: str = (
         'question: \n```\n{query}\n```\ncontext: ```\n{context}\n```'
     )
-    structured_output: Optional[BaseModel] = None
+    structured_output: Optional[Type[BaseModel]] = None
 
     # default parameters that can be overwritten by the derived class
     default_device_name: str = 'cpu'
@@ -46,7 +46,7 @@ def __init__(
         prompt_template: str = '',
         output_max_length: int = 500,
         device: str = 'auto',
-        structured_output: Optional[BaseModel] = None,
+        structured_output: Optional[Type[BaseModel]] = None,
         logs: dict[str, Any] = {},
     ) -> None:
         """Initialize Generation class.
@@ -61,7 +61,7 @@ def __init__(
         output_max_length : int
             Maximum length of the generated output.
         device: str (default=auto)
-        structured_output: Optional[BaseModel] = None
+        structured_output: Optional[Type[BaseModel]] = None
         logs: dict[str, Any] = {}
         """
         self.api_key: str = api_key
@@ -74,7 +74,7 @@ def __init__(
         self.prompt_template: str = (
             prompt_template or self.default_prompt_template
         )
-        self.structured_output: Optional[BaseModel] = None
+        self.structured_output: Optional[Type[BaseModel]] = structured_output
 
         if device not in ['cpu', 'cuda', 'auto']:
             raise Exception(
@@ -105,7 +105,7 @@ def generate(
         self,
         query: str,
         context: list[str],
-    ) -> str:
+    ) -> str | BaseModel:
         """Generate text with optional language parameter.
 
         Parameters

src/rago/generation/gemini.py

@@ -7,6 +7,7 @@
 import google.generativeai as genai
 import instructor
 
+from pydantic import BaseModel
 from typeguard import typechecked
 
 from rago.generation.base import GenerationBase
@@ -24,20 +25,38 @@ def _setup(self) -> None:
         model = genai.GenerativeModel(self.model_name)
 
         self.model = (
-            instructor.from_gemini(model) if self.structured_output else model
+            instructor.from_gemini(
+                client=model,
+                mode=instructor.Mode.GEMINI_JSON,
+            )
+            if self.structured_output
+            else model
         )
 
-    def generate(self, query: str, context: list[str]) -> str:
+    def generate(self, query: str, context: list[str]) -> str | BaseModel:
         """Generate text using Gemini model support."""
         input_text = self.prompt_template.format(
             query=query, context=' '.join(context)
         )
 
+        if not self.structured_output:
+            models_params_gen = {'contents': input_text}
+            response = self.model.generate_content(**models_params_gen)
+            self.logs['model_params'] = models_params_gen
+            return cast(str, response.text.strip())
+
+        messages = [
+            {'role': 'user', 'content': input_text},
+        ]
         model_params = {
-            'contents': input_text,
+            'messages': messages,
+            'response_model': self.structured_output,
         }
 
-        response = self.model.generate_content(**model_params)
+        response = self.model.create(
+            **model_params,
+        )
 
         self.logs['model_params'] = model_params
-        return cast(str, response.text.strip())
+
+        return cast(BaseModel, response)

src/rago/generation/openai.py

@@ -7,6 +7,7 @@
 import instructor
 import openai
 
+from pydantic import BaseModel
 from typeguard import typechecked
 
 from rago.generation.base import GenerationBase
@@ -30,7 +31,7 @@ def generate(
         self,
         query: str,
         context: list[str],
-    ) -> str:
+    ) -> str | BaseModel:
         """Generate text using OpenAI's API with dynamic model support."""
         input_text = self.prompt_template.format(
             query=query, context=' '.join(context)
@@ -49,8 +50,15 @@ def generate(
             presence_penalty=0.3,
         )
 
+        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(str, response.choices[0].message.content.strip())
+        has_choices = hasattr(response, 'choices')
+
+        if has_choices and isinstance(response.choices, list):
+            return cast(str, response.choices[0].message.content.strip())
+        return cast(BaseModel, response)

tests/models.py

@@ -0,0 +1,28 @@
+"""Models used for the unit tests."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class AnimalModel(BaseModel):
+    """Model for animals."""
+
+    name: Literal[
+        'Blue Whale',
+        'Peregrine Falcon',
+        'Giant Panda',
+        'Cheetah',
+        'Komodo Dragon',
+        'Arctic Fox',
+        'Monarch Butterfly',
+        'Great White Shark',
+        'Honey Bee',
+        'Emperor Penguin',
+        'Unknown',
+    ] = Field(
+        ...,
+        description='The predicted class label.',
+    )

tests/test_gemini.py

@@ -2,13 +2,17 @@
 
 import os
 
+from typing import cast
+
 import pytest
 
 from rago import Rago
 from rago.augmented import SentenceTransformerAug
 from rago.generation import GeminiGen
 from rago.retrieval import StringRet
 
+from .models import AnimalModel
+
 
 @pytest.fixture
 def api_key(env) -> str:
@@ -52,3 +56,53 @@ def test_gemini_generation(animals_data: list[str], api_key: str) -> None:
     assert logs['retrieval']
     assert logs['augmented']
     assert logs['generation']
+
+
+@pytest.mark.skip_on_ci
+@pytest.mark.parametrize(
+    'question,expected_answer',
+    [
+        ('What animal is larger than a dinosaur?', 'Blue Whale'),
+        (
+            'What animal is renowned as the fastest animal on the planet?',
+            'Peregrine Falcon',
+        ),
+    ],
+)
+def test_rag_gemini_structured_output(
+    api_key: str,
+    animals_data: list[str],
+    question: str,
+    expected_answer: str,
+) -> None:
+    """Test RAG pipeline with Gemini."""
+    logs = {
+        'retrieval': {},
+        'augmented': {},
+        'generation': {},
+    }
+
+    rag = Rago(
+        retrieval=StringRet(animals_data, logs=logs['retrieval']),
+        augmented=SentenceTransformerAug(top_k=3, logs=logs['augmented']),
+        generation=GeminiGen(
+            api_key=api_key,
+            model_name='gemini-1.5-flash',
+            logs=logs['generation'],
+            structured_output=AnimalModel,
+        ),
+    )
+
+    result = cast(AnimalModel, rag.prompt(question))
+
+    error_message = (
+        f'Expected response to mention `{expected_answer}`. '
+        f'Result: `{result.name}`.'
+    )
+
+    assert expected_answer == result.name, error_message
+
+    # check if logs have been used
+    assert logs['retrieval']
+    assert logs['augmented']
+    assert logs['generation']

tests/test_openai.py

@@ -2,13 +2,17 @@
 
 import os
 
+from typing import cast
+
 import pytest
 
 from rago import Rago
 from rago.augmented import OpenAIAug
 from rago.generation import OpenAIGen
 from rago.retrieval import StringRet
 
+from .models import AnimalModel
+
 
 @pytest.fixture
 def api_key(env) -> str:
@@ -82,3 +86,53 @@ def test_rag_openai_gpt(animals_data: list[str], api_key: str) -> None:
     assert logs['retrieval']
     assert logs['augmented']
     assert logs['generation']
+
+
+@pytest.mark.skip_on_ci
+@pytest.mark.parametrize(
+    'question,expected_answer',
+    [
+        ('What animal is larger than a dinosaur?', 'Blue Whale'),
+        (
+            'What animal is renowned as the fastest animal on the planet?',
+            'Peregrine Falcon',
+        ),
+    ],
+)
+def test_rag_openai_gpt_structured_output(
+    api_key: str,
+    animals_data: list[str],
+    question: str,
+    expected_answer: str,
+) -> None:
+    """Test RAG pipeline with OpenAI's GPT."""
+    logs = {
+        'retrieval': {},
+        'augmented': {},
+        'generation': {},
+    }
+
+    rag = Rago(
+        retrieval=StringRet(animals_data, logs=logs['retrieval']),
+        augmented=OpenAIAug(api_key=api_key, top_k=3, logs=logs['augmented']),
+        generation=OpenAIGen(
+            api_key=api_key,
+            model_name='gpt-3.5-turbo',
+            logs=logs['generation'],
+            structured_output=AnimalModel,
+        ),
+    )
+
+    result = cast(AnimalModel, rag.prompt(question))
+
+    error_message = (
+        f'Expected response to mention `{expected_answer}`. '
+        f'Result: `{result.name}`.'
+    )
+
+    assert expected_answer == result.name, error_message
+
+    # check if logs have been used
+    assert logs['retrieval']
+    assert logs['augmented']
+    assert logs['generation']