Update rag_agent.py

rag_agent.py

@@ -1,24 +1,38 @@
 import os
 from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, ServiceContext
-from llama_index.llms.openai import OpenAI
+from llama_index.llms.huggingface import HuggingFaceLLM
 from llama_index.core.tools import QueryEngineTool, ToolMetadata
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Build Service Context with preferred model
-llm = OpenAI(model="gpt-3.5-turbo", temperature=0)
+# Load local or remote HF model (example: mistral hosted model)
+model_name = "mistralai/Mistral-7B-Instruct-v0.2"
+
+# Create HuggingFaceLLM
+llm = HuggingFaceLLM(
+    model_name=model_name,
+    tokenizer_name=model_name,
+    context_window=2048,
+    max_new_tokens=512,
+    generate_kwargs={"temperature": 0.1},
+    tokenizer_kwargs={"padding_side": "left"},
+    device_map="auto",
+)
+
+# Build service context
 service_context = ServiceContext.from_defaults(llm=llm)
 
-# Load and index documents
+# Load knowledge base
 documents = SimpleDirectoryReader("kb").load_data()
 index = VectorStoreIndex.from_documents(documents, service_context=service_context)
 query_engine = index.as_query_engine()
 
-# Tool wrapper to integrate with a multi-tool agent if needed
+# Wrap in a tool
 rag_tool = QueryEngineTool(
     query_engine=query_engine,
-    metadata=ToolMetadata(name="RAGSearch", description="Answers questions using a local knowledge base")
+    metadata=ToolMetadata(name="RAGSearch", description="Answers from a local HF-based RAG system.")
 )
 
-# Agent class (Hugging Face-compatible)
+# Agent definition
 class BasicAgent:
     def __init__(self):
         self.tool = rag_tool