-
Notifications
You must be signed in to change notification settings - Fork 682
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
5 changed files
with
208 additions
and
14 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,166 @@ | ||
| import asyncio | ||
| import os | ||
| from collections.abc import Iterator | ||
| import vertexai.preview | ||
| from vertexai.preview.generative_models import GenerativeModel, Content, Part | ||
| from wasabi import msg | ||
|
|
||
| from dotenv import load_dotenv | ||
|
|
||
| from goldenverba.components.generation.interface import Generator | ||
|
|
||
| load_dotenv() | ||
|
|
||
|
|
||
| class GeminiGenerator(Generator): | ||
| """ | ||
| Gemini Generator. | ||
| """ | ||
|
|
||
| def __init__(self): | ||
| super().__init__() | ||
| self.name = "GeminiGenerator" | ||
| self.description = "Generator using Google's Gemini 1.5 Pro model" | ||
| self.requires_library = ["google-cloud-aiplatform"] | ||
| self.requires_env = [ | ||
| "GOOGLE_APPLICATION_CREDENTIALS", | ||
| "GOOGLE_CLOUD_PROJECT", | ||
| ] | ||
| self.streamable = True | ||
| self.model_name = os.getenv("GEMINI_MODEL", "gemini-1.5-pro-preview-0409") | ||
| self.context_window = 10000 | ||
| self.system_instruction = "You are Verba, The Golden RAGtriever, a chatbot for Retrieval Augmented Generation (RAG). You will receive a user query and context pieces that have a semantic similarity to that specific query. Please answer these user queries only their provided context. If the provided documentation does not provide enough information, say so. If the user asks questions about you as a chatbot specifially, answer them naturally. If the answer requires code examples encapsulate them with ```programming-language-name ```. Don't do pseudo-code." | ||
|
|
||
| async def generate( | ||
| self, | ||
| queries: list[str], | ||
| context: list[str], | ||
| conversation: dict = None, | ||
| ) -> str: | ||
| """Generate an answer based on a list of queries and list of contexts, and includes conversational context | ||
| @parameter: queries : list[str] - List of queries | ||
| @parameter: context : list[str] - List of contexts | ||
| @parameter: conversation : dict - Conversational context | ||
| @returns str - Answer generated by the Generator. | ||
| """ | ||
| if conversation is None: | ||
| conversation = {} | ||
| messages = self.prepare_messages(queries, context, conversation) | ||
|
|
||
| try: | ||
|
|
||
| project_id = os.getenv("GOOGLE_CLOUD_PROJECT") | ||
|
|
||
| REGION = "us-central1" | ||
| vertexai.init(project=project_id, location=REGION) | ||
|
|
||
| generative_multimodal_model = GenerativeModel( | ||
| "gemini-1.5-pro-preview-0409", | ||
| system_instruction=self.system_instruction, | ||
| ) | ||
| # response = generative_multimodal_model.generate_content(messages) | ||
|
|
||
| completion = await asyncio.to_thread( | ||
| generative_multimodal_model.generate_content, **messages | ||
| ) | ||
| system_msg = str(completion.candidates[0].content.parts[0].text) | ||
|
|
||
| except Exception: | ||
| raise | ||
|
|
||
| return system_msg | ||
|
|
||
| async def generate_stream( | ||
| self, | ||
| queries: list[str], | ||
| context: list[str], | ||
| conversation: dict = None, | ||
| ) -> Iterator[dict]: | ||
| """Generate a stream of response dicts based on a list of queries and list of contexts, and includes conversational context | ||
| @parameter: queries : list[str] - List of queries | ||
| @parameter: context : list[str] - List of contexts | ||
| @parameter: conversation : dict - Conversational context | ||
| @returns Iterator[dict] - Token response generated by the Generator in this format {system:TOKEN, finish_reason:stop or empty}. | ||
| """ | ||
| if conversation is None: | ||
| conversation = {} | ||
| messages = self.prepare_messages(queries, context, conversation) | ||
|
|
||
| try: | ||
|
|
||
| project_id = os.getenv("GOOGLE_CLOUD_PROJECT") | ||
|
|
||
| REGION = "us-central1" | ||
| vertexai.init(project=project_id, location=REGION) | ||
|
|
||
| generative_multimodal_model = GenerativeModel( | ||
| "gemini-1.5-pro-preview-0409", | ||
| system_instruction=self.system_instruction, | ||
| ) | ||
| # response = generative_multimodal_model.generate_content(messages) | ||
|
|
||
| completion = await generative_multimodal_model.generate_content_async( | ||
| stream=True, contents=messages | ||
| ) | ||
|
|
||
| print("using google gemini") | ||
|
|
||
| iter = completion.__aiter__() | ||
|
|
||
| try: | ||
| while True: | ||
| chunk = await iter.__anext__() | ||
| if len(chunk.candidates) > 0: | ||
| msg.info(chunk.candidates[0]) | ||
| if len(chunk.candidates[0].content.parts[0].text) > 0: | ||
| yield { | ||
| "message": chunk.candidates[0].content.parts[0].text, | ||
| "finish_reason": chunk.candidates[0].finish_reason, | ||
| } | ||
| else: | ||
| yield { | ||
| "message": "", | ||
| "finish_reason": chunk.candidates[0].finish_reason, | ||
| } | ||
| except StopAsyncIteration: | ||
| pass | ||
|
|
||
| except Exception: | ||
| raise | ||
|
|
||
| def prepare_messages( | ||
| self, queries: list[str], context: list[str], conversation: dict[str, str] | ||
| ) -> list[Content]: | ||
| """ | ||
| Prepares a list of messages formatted for a Retrieval Augmented Generation chatbot system, including system instructions, previous conversation, and a new user query with context. | ||
| @parameter queries: A list of strings representing the user queries to be answered. | ||
| @parameter context: A list of strings representing the context information provided for the queries. | ||
| @parameter conversation: A list of previous conversation messages that include the role and content. | ||
| @returns A list of message dictionaries formatted for the chatbot. This includes an initial system message, the previous conversation messages, and the new user query encapsulated with the provided context. | ||
| Each message in the list is a dictionary with 'role' and 'content' keys, where 'role' is either 'system' or 'user', and 'content' contains the relevant text. This will depend on the LLM used. | ||
| """ | ||
| messages = [] | ||
|
|
||
| for message in conversation: | ||
| messages.append( | ||
| Content(role=message.type, parts=[Part.from_text(message.content)]) | ||
| ) | ||
|
|
||
| query = " ".join(queries) | ||
| user_context = " ".join(context) | ||
|
|
||
| messages.append( | ||
| Content( | ||
| role="user", | ||
| parts=[ | ||
| Part.from_text( | ||
| f"Please answer this query: '{query}' with this provided context: {user_context}" | ||
| ) | ||
| ], | ||
| ) | ||
| ) | ||
|
|
||
| return messages |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters