| 121 | |
| 122 | |
| 123 | def init_gemini(): |
| 124 | from llama_index.llms.gemini import Gemini |
| 125 | from llama_index.embeddings.gemini import GeminiEmbedding |
| 126 | |
| 127 | model_map: Dict[str, str] = { |
| 128 | "gemini-1.5-pro-latest": "models/gemini-1.5-pro-latest", |
| 129 | "gemini-pro": "models/gemini-pro", |
| 130 | "gemini-pro-vision": "models/gemini-pro-vision", |
| 131 | } |
| 132 | |
| 133 | embed_model_map: Dict[str, str] = { |
| 134 | "embedding-001": "models/embedding-001", |
| 135 | "text-embedding-004": "models/text-embedding-004", |
| 136 | } |
| 137 | |
| 138 | Settings.llm = Gemini(model=model_map[os.getenv("MODEL")]) |
| 139 | Settings.embed_model = GeminiEmbedding( |
| 140 | model_name=embed_model_map[os.getenv("EMBEDDING_MODEL")] |
| 141 | ) |