LlamaIndex: Make model configurable

Author: amotl

topic/machine-learning/llama-index/boot.py

@@ -1,15 +1,20 @@
 import os
+from typing import Tuple
 
 import openai
 from langchain_openai import AzureOpenAIEmbeddings
 from langchain_openai import OpenAIEmbeddings
-from llama_index.core import Settings
+from llama_index.core.base.embeddings.base import BaseEmbedding
+from llama_index.core.llms import LLM
 from llama_index.llms.azure_openai import AzureOpenAI
 from llama_index.llms.openai import OpenAI
 from llama_index.embeddings.langchain import LangchainEmbedding
 
 
-def configure_llm():
+MODEL_NAME = "gpt-4o"
+
+
+def configure_llm() -> Tuple[LLM, BaseEmbedding]:
     """
     Configure LLM. Use either vanilla Open AI, or Azure Open AI.
     """
@@ -21,27 +26,32 @@ def configure_llm():
 
     if openai.api_type == "openai":
         llm = OpenAI(
+            model=MODEL_NAME,
+            temperature=0.0,
             api_key=os.getenv("OPENAI_API_KEY"),
-            temperature=0.0
         )
     elif openai.api_type == "azure":
         llm = AzureOpenAI(
+            model=MODEL_NAME,
+            temperature=0.0,
             engine=os.getenv("LLM_INSTANCE"),
             azure_endpoint=os.getenv("OPENAI_AZURE_ENDPOINT"),
             api_key = os.getenv("OPENAI_API_KEY"),
             api_version = os.getenv("OPENAI_AZURE_API_VERSION"),
-            temperature=0.0
         )
     else:
         raise ValueError(f"Open AI API type not defined or invalid: {openai.api_type}")
 
-    Settings.llm = llm
     if openai.api_type == "openai":
-        Settings.embed_model = LangchainEmbedding(OpenAIEmbeddings())
+        embed_model = LangchainEmbedding(OpenAIEmbeddings(model=MODEL_NAME))
     elif openai.api_type == "azure":
-        Settings.embed_model = LangchainEmbedding(
+        embed_model = LangchainEmbedding(
             AzureOpenAIEmbeddings(
                 azure_endpoint=os.getenv("OPENAI_AZURE_ENDPOINT"),
                 model=os.getenv("EMBEDDING_MODEL_INSTANCE")
             )
         )
+    else:
+        embed_model = None
+
+    return llm, embed_model

topic/machine-learning/llama-index/demo_mcp.py

@@ -24,14 +24,17 @@
 
 from dotenv import load_dotenv
 from llama_index.core.agent.workflow import FunctionAgent
-from llama_index.llms.openai import OpenAI
+from llama_index.core.llms import LLM
 from llama_index.tools.mcp import BasicMCPClient, McpToolSpec
 
 from boot import configure_llm
 
 
 class Agent:
 
+    def __init__(self, llm: LLM):
+        self.llm = llm
+
     async def get_tools(self):
         # Connect to the CrateDB MCP server using `streamable-http` transport.
         mcp_url = os.getenv("CRATEDB_MCP_URL", "http://127.0.0.1:8000/mcp/")
@@ -49,7 +52,7 @@ async def get_agent(self):
         return FunctionAgent(
             name="Agent",
             description="CrateDB text-to-SQL agent",
-            llm=OpenAI(model="gpt-4o"),
+            llm=self.llm,
             tools=await self.get_tools(),
             system_prompt=Instructions.full(),
         )
@@ -69,10 +72,10 @@ def main():
 
     # Configure application.
     load_dotenv()
-    configure_llm()
+    llm, embed_model = configure_llm()
 
     # Use an agent that uses the CrateDB MCP server.
-    agent = Agent()
+    agent = Agent(llm)
 
     # Invoke an inquiry.
     print("Running query")

topic/machine-learning/llama-index/demo_nlsql.py

@@ -9,7 +9,6 @@
 from dotenv import load_dotenv
 from llama_index.core.utilities.sql_wrapper import SQLDatabase
 from llama_index.core.query_engine import NLSQLTableQueryEngine
-from llama_index.core import Settings
 
 from boot import configure_llm
 
@@ -21,7 +20,7 @@ def main():
 
     # Configure application.
     load_dotenv()
-    configure_llm()
+    llm, embed_model = configure_llm()
 
     # Configure database connection and query engine.
     print("Connecting to CrateDB")
@@ -33,7 +32,8 @@ def main():
     query_engine = NLSQLTableQueryEngine(
         sql_database=sql_database,
         tables=[os.getenv("CRATEDB_TABLE_NAME")],
-        llm=Settings.llm
+        llm=llm,
+        embed_model=embed_model,
     )
 
     # Invoke an inquiry.