Update agent.py

agent.py

@@ -33,6 +33,7 @@ from llama_index.readers.youtube_transcript import YoutubeTranscriptReader
 from llama_index.tools.arxiv import ArxivToolSpec
 from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec
 from llama_index.core.agent.workflow import AgentWorkflow
+from llama_index.llms.vllm import Vllm
 
 # Import all required official LlamaIndex Readers
 from llama_index.readers.file import (
@@ -48,7 +49,6 @@ try:
     # Gemini (for API mode)
     from llama_index.llms.gemini import Gemini
     from llama_index.embeddings.gemini import GeminiEmbedding
-    from llama_index_llms_vllm import Vllm
     GEMINI_AVAILABLE = True
 except ImportError:
     GEMINI_AVAILABLE = False
@@ -120,36 +120,31 @@ def initialize_models(use_api_mode=False):
             # Try to use Pixtral 12B with vLLM if available
             pixtral_model = "Qwen/Qwen3-8B"  # Fallback model
             try:
-                if importlib.util.find_spec("llama_index_llms_vllm") is not None:
-                    from llama_index_llms_vllm import Vllm
-                    # Check if Pixtral 12B is accessible
-                    if os.path.exists("/path/to/pixtral-12b") or True:  # Placeholder check
-                        pixtral_model = "mistralai/pixtral-12b"
-                        print(f"Using Pixtral 12B with vLLM")
-
-                        # Custom prompt template for Pixtral model
-                        def messages_to_prompt(messages):
-                            prompt = "\n".join([str(x) for x in messages])
-                            return f"<s>[INST] {prompt} [/INST] </s>\n"
-
-                        def completion_to_prompt(completion):
-                            return f"<s>[INST] {completion} [/INST] </s>\n"
-
-                        proj_llm = Vllm(
-                            model=pixtral_model,
-                            tensor_parallel_size=1,  # Adjust based on available GPUs
-                            max_new_tokens=16000,
-                            messages_to_prompt=messages_to_prompt,
-                            completion_to_prompt=completion_to_prompt,
-                            temperature=0.6,
-                            top_p=0.95,
-                            top_k=20
-                        )
-                    else:
-                        # Use regular Qwen model if Pixtral not found
-                        raise ImportError("Pixtral 12B not found")
+                if os.path.exists("/path/to/pixtral-12b") or True:  # Placeholder check
+                    pixtral_model = "mistralai/pixtral-12b"
+                    print(f"Using Pixtral 12B with vLLM")
+
+                    # Custom prompt template for Pixtral model
+                    def messages_to_prompt(messages):
+                        prompt = "\n".join([str(x) for x in messages])
+                        return f"<s>[INST] {prompt} [/INST] </s>\n"
+
+                    def completion_to_prompt(completion):
+                        return f"<s>[INST] {completion} [/INST] </s>\n"
+
+                    proj_llm = Vllm(
+                        model=pixtral_model,
+                        tensor_parallel_size=1,  # Adjust based on available GPUs
+                        max_new_tokens=16000,
+                        messages_to_prompt=messages_to_prompt,
+                        completion_to_prompt=completion_to_prompt,
+                        temperature=0.6,
+                        top_p=0.95,
+                        top_k=20
+                    )
                 else:
-                    raise ImportError("vLLM not available")
+                    # Use regular Qwen model if Pixtral not found
+                    raise ImportError("Pixtral 12B not found")
             except (ImportError, Exception) as e:
                 print(f"Error loading Pixtral with vLLM: {e}")
                 print(f"Falling back to {pixtral_model} with HuggingFace...")