feat: enhanced model downloads

Author: jasperan

agentic_rag/gradio_app.py

@@ -103,7 +103,7 @@ def chat(message: str, history: List[List[str]], agent_type: str, use_cot: bool,
             model_type = "Local (Mistral)"
         elif "Ollama" in agent_type:
             model_type = "Ollama"
-            # Extract model name from agent_type
+            # Extract model name from agent_type and use correct Ollama model names
             if "llama3" in agent_type.lower():
                 model_name = "ollama:llama3"
             elif "phi-3" in agent_type.lower():
@@ -588,6 +588,18 @@ def download_model(model_type: str) -> str:
                 print(f"Pulling Ollama model: {model_name}")
                 start_time = time.time()
 
+                # Check if model already exists
+                try:
+                    models = ollama.list().models
+                    available_models = [model.model for model in models]
+                    
+                    # Check for model with or without :latest suffix
+                    if model_name in available_models or f"{model_name}:latest" in available_models:
+                        return f"✅ Model {model_name} is already available in Ollama."
+                except Exception:
+                    # If we can't check, proceed with pull anyway
+                    pass
+                
                 # Pull the model with progress tracking
                 progress_text = ""
                 for progress in ollama.pull(model_name, stream=True):

agentic_rag/local_rag_agent.py

@@ -51,7 +51,8 @@ def __init__(self, model_name: str):
         Args:
             model_name: Name of the Ollama model to use
         """
-        self.model_name = model_name
+        # Remove the 'ollama:' prefix if present
+        self.model_name = model_name.replace("ollama:", "") if model_name.startswith("ollama:") else model_name
         self._check_ollama_running()
 
     def _check_ollama_running(self):
@@ -67,8 +68,13 @@ def _check_ollama_running(self):
 
                 # Check if the requested model is available
                 if self.model_name not in available_models:
-                    print(f"Model '{self.model_name}' not found in Ollama. Available models: {', '.join(available_models)}")
-                    print(f"You can pull it with: ollama pull {self.model_name}")
+                    # Try with :latest suffix
+                    if f"{self.model_name}:latest" in available_models:
+                        self.model_name = f"{self.model_name}:latest"
+                        print(f"Using model with :latest suffix: {self.model_name}")
+                    else:
+                        print(f"Model '{self.model_name}' not found in Ollama. Available models: {', '.join(available_models)}")
+                        print(f"You can pull it with: ollama pull {self.model_name}")
             except Exception as e:
                 raise ConnectionError(f"Failed to connect to Ollama. Please make sure Ollama is running. Error: {str(e)}")
 
@@ -427,6 +433,33 @@ def _generate_response(self, query: str, context: List[Dict[str, Any]]) -> Dict[
         prompt = template.format(context=context_str, query=query)
         response = self._generate_text(prompt)
 
+        # Add sources to response if available
+        if context:
+            # Group sources by document
+            sources = {}
+            for item in context:
+                source = item['metadata'].get('source', 'Unknown')
+                if source not in sources:
+                    sources[source] = set()
+                
+                # Add page number if available
+                if 'page' in item['metadata']:
+                    sources[source].add(str(item['metadata']['page']))
+                # Add file path if available for code
+                if 'file_path' in item['metadata']:
+                    sources[source] = item['metadata']['file_path']
+            
+            # Print concise source information
+            print("\nSources detected:")
+            for source, details in sources.items():
+                if isinstance(details, set):  # PDF with pages
+                    pages = ", ".join(sorted(details))
+                    print(f"Document: {source} (pages: {pages})")
+                else:  # Code with file path
+                    print(f"Code file: {source}")
+            
+            response['sources'] = sources
+        
         return {
             "answer": response,
             "context": context

agentic_rag/rag_agent.py

@@ -190,25 +190,51 @@ def _process_query_standard(self, query: str) -> Dict[str, Any]:
         return response
 
     def _generate_response(self, query: str, context: List[Dict[str, Any]]) -> Dict[str, Any]:
-        """Generate a response using the retrieved context"""
-        context_str = "\n\n".join([f"Context {i+1}:\n{item['content']}" 
-                                  for i, item in enumerate(context)])
+        """Generate a response based on the query and context"""
+        # Format context for the prompt
+        formatted_context = "\n\n".join([f"Context {i+1}:\n{item['content']}" 
+                                       for i, item in enumerate(context)])
 
-        template = """Answer the following query using the provided context. 
-Respond as if you are knowledgeable about the topic and incorporate the context naturally.
-Do not mention limitations in the context or that you couldn't find specific information.
-
-Context:
-{context}
-
-Query: {query}
-
-Answer:"""
+        # Create the prompt
+        system_prompt = """You are an AI assistant answering questions based on the provided context.
+Answer the question based on the context provided. If the answer is not in the context, say "I don't have enough information to answer this question." Be concise and accurate."""
 
-        prompt = ChatPromptTemplate.from_template(template)
-        messages = prompt.format_messages(context=context_str, query=query)
+        # Create messages for the chat model
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": f"Context:\n{formatted_context}\n\nQuestion: {query}"}
+        ]
+        
+        # Generate response
         response = self.llm.invoke(messages)
 
+        # Add sources to response if available
+        if context:
+            # Group sources by document
+            sources = {}
+            for item in context:
+                source = item['metadata'].get('source', 'Unknown')
+                if source not in sources:
+                    sources[source] = set()
+                
+                # Add page number if available
+                if 'page' in item['metadata']:
+                    sources[source].add(str(item['metadata']['page']))
+                # Add file path if available for code
+                if 'file_path' in item['metadata']:
+                    sources[source] = item['metadata']['file_path']
+            
+            # Print concise source information
+            print("\nSources detected:")
+            for source, details in sources.items():
+                if isinstance(details, set):  # PDF with pages
+                    pages = ", ".join(sorted(details))
+                    print(f"Document: {source} (pages: {pages})")
+                else:  # Code with file path
+                    print(f"Code file: {source}")
+            
+            response['sources'] = sources
+        
         return {
             "answer": response.content,
             "context": context