Spaces:
Running
Running
| from typing import List, Literal, cast | |
| from pydantic import SecretStr | |
| from _utils.google_integration.google_cloud import GCP_PROJECT, upload_to_gcs | |
| from setup.easy_imports import ChatOpenAI, ChatGoogleGenerativeAI | |
| import os | |
| from langchain_core.messages import HumanMessage | |
| from langchain_google_vertexai import ChatVertexAI | |
| deepseek_api_key = cast(str, os.environ.get("DEEPSEEKK_API_KEY")) | |
| google_api_key = cast(str, os.environ.get("GOOGLE_API_KEY_PEIXE")) | |
| open_ai_token = cast(str, os.environ.get("OPENAI_API_KEY")) | |
| Google_llms = Literal[ | |
| "gemini-2.5-pro-preview-05-06", | |
| "gemini-2.0-flash", | |
| "gemini-2.0-flash-lite", | |
| "gemini-2.5-flash-preview-04-17", | |
| ] | |
| class LLM: | |
| def __init__(self): | |
| pass | |
| def open_ai(self, model="gpt-4o-mini"): | |
| return ChatOpenAI(api_key=SecretStr(open_ai_token), model=model) | |
| def deepseek(self, model="deepseek-chat"): | |
| return ChatOpenAI( | |
| api_key=SecretStr(deepseek_api_key), | |
| base_url="https://api.deepseek.com/v1", | |
| model=model, | |
| ) | |
| def google_gemini( | |
| self, | |
| model: Google_llms = "gemini-2.0-flash", | |
| ): | |
| return ChatGoogleGenerativeAI( | |
| api_key=SecretStr(google_api_key), | |
| model=model, | |
| temperature=0, | |
| max_tokens=None, | |
| timeout=None, | |
| max_retries=2, | |
| ) | |
| async def google_gemini_ainvoke( | |
| self, | |
| prompt: str, | |
| model: Google_llms = "gemini-2.0-flash", | |
| max_retries: int = 3, | |
| ): | |
| for attempt in range(max_retries): | |
| try: | |
| response = await self.google_gemini(model).ainvoke( | |
| [HumanMessage(content=prompt)] | |
| ) | |
| if isinstance(response.content, list): | |
| response.content = "\n".join(response.content) # type: ignore | |
| return response | |
| except Exception as e: | |
| model = "gemini-2.0-flash" | |
| print(f"Attempt {attempt + 1} failed with error: {e}") | |
| # Final attempt fallback logic (optional) | |
| try: | |
| print("Final attempt with fallback model...") | |
| response = await self.open_ai("chat-gpt-4o-mini").ainvoke( | |
| [HumanMessage(content=prompt)] | |
| ) | |
| return response | |
| except Exception as e: | |
| raise Exception( | |
| "Failed to generate the final document after 5 retries and the fallback attempt with chat-gpt-4o-mini." | |
| ) from e | |
| async def google_gemini_vertex_ainvoke( | |
| self, | |
| prompt: str, | |
| list_of_pdfs: List[str], | |
| model: Google_llms = "gemini-2.5-flash-preview-04-17", | |
| max_retries: int = 3, | |
| ) -> str | None: | |
| message_parts = [ | |
| {"type": "text", "text": prompt}, | |
| ] | |
| for pdf in list_of_pdfs: | |
| pdf_gcs_uri = upload_to_gcs(pdf) | |
| message_parts.append( | |
| { | |
| # This structure is used for file references via URI | |
| "type": "media", | |
| "mime_type": "application/pdf", # <-- mime_type moved up | |
| "file_uri": pdf_gcs_uri, # <-- file_uri moved up | |
| } | |
| ) | |
| for attempt in range(max_retries): | |
| try: | |
| llm = ChatVertexAI( | |
| model_name=model, | |
| project=GCP_PROJECT, | |
| location="us-central1", | |
| temperature=0, | |
| ) | |
| response = await llm.ainvoke( | |
| [HumanMessage(content=message_parts)] # type: ignore | |
| ) | |
| if isinstance(response.content, list): | |
| response.content = "\n".join(response.content) # type: ignore | |
| return response.content # type: ignore | |
| except Exception as e: | |
| model = "gemini-2.0-flash" | |
| print(f"Attempt {attempt + 1} failed with error: {e}") | |