Update app/core/llm_engine.py

app/core/llm_engine.py

@@ -1,41 +1,107 @@
 # # llm_engine.py
 
-import google.generativeai as genai
-from app.core.config import GEMINI_API_KEY
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_nvidia_ai_endpoints import ChatNVIDIA
-import os
+# from app.core.config import NVIDIA_API_KEY
+# from langchain_nvidia_ai_endpoints import ChatNVIDIA
+# import os
 
-# ✅ Configure Gemini client
-genai.configure(api_key=GEMINI_API_KEY)
 
-# llm = ChatGoogleGenerativeAI(
-#     model="gemini-2.5-flash",
-#     google_api_key=GEMINI_API_KEY,
-#     temperature=0.2,
-#     max_output_tokens=800,
+# llm = ChatNVIDIA(
+#     model="meta/llama-3.1-70b-instruct",
+#     api_key=os.getenv("NVIDIA_API_KEY"),
+#     temperature=0.7,
+#     max_tokens=1024
+# )
+
+# eval_llm = ChatNVIDIA(
+#     # Faster for evaluation
+#     model="meta/llama-3.1-8b-instruct", 
+#     temperature=0.0,
+#     max_tokens=200
 # )
 
 
+
+
+# app/core/llm_engine.py
+
+import os
+
+from langchain_nvidia_ai_endpoints import ChatNVIDIA
+
+# ============================================================
+# Configuration
+# ============================================================
+
+NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY")
+
+MAIN_MODEL = "meta/llama-3.1-70b-instruct"
+EVAL_MODEL = "meta/llama-3.1-8b-instruct"
+
+
+# ============================================================
+# Main LLM (Non-streaming)
+# Used everywhere graph.invoke() is still called.
+# ============================================================
+
 llm = ChatNVIDIA(
-    model="meta/llama-3.1-70b-instruct",  # or nvidia/nemotron-4-340b-instruct
-    api_key=os.getenv("NVIDIA_API_KEY"),
+    model=MAIN_MODEL,
+    api_key=NVIDIA_API_KEY,
     temperature=0.7,
-    max_tokens=1024
+    max_tokens=1024,
 )
 
+# ============================================================
+# Streaming LLM
+# Used by /query-stream endpoint.
+# Supports .astream()
+# ============================================================
+
+streaming_llm = ChatNVIDIA(
+    model=MAIN_MODEL,
+    api_key=NVIDIA_API_KEY,
+    temperature=0.7,
+    max_tokens=1024,
+    streaming=True,
+)
+
+# ============================================================
+# Evaluator LLM
+# Faster + deterministic
+# ============================================================
+
 eval_llm = ChatNVIDIA(
-    model="meta/llama-3.1-8b-instruct",  # Faster for evaluation
+    model=EVAL_MODEL,
+    api_key=NVIDIA_API_KEY,
     temperature=0.0,
-    max_tokens=200
+    max_tokens=200,
 )
 
+# ============================================================
+# Helper getters
+# (optional, but keeps imports clean)
+# ============================================================
 
-# Separate LLM for evaluator — needs near-deterministic JSON output
+def get_llm():
+    """
+    Standard synchronous LLM.
+    """
+    return llm
+
+
+def get_streaming_llm():
+    """
+    Streaming LLM.
+    Use with:
+
+        async for chunk in llm.astream(...):
+            ...
+    """
+    return streaming_llm
+
+
+def get_eval_llm():
+    """
+    Evaluator model.
+    """
+    return eval_llm
 
-# eval_llm = ChatGoogleGenerativeAI(
-#     model="gemini-2.0-flash",
-#     google_api_key=GEMINI_API_KEY,
-#     temperature=0.0,
-#     max_output_tokens=200,
-# )