huggingface model support draft

src/know_lang_bot/chat_bot/__main__.py

@@ -19,7 +19,6 @@ async def test_chat_processing():
     )
     
     print(f"Answer: {result.answer}")
-    print(f"References: {result.references_md}")
 
 if __name__ == "__main__":
     asyncio.run(test_chat_processing())

src/know_lang_bot/chat_bot/chat_graph.py

@@ -12,8 +12,11 @@ from pydantic_ai import Agent
 import logfire
 from pprint import pformat
 from enum import Enum
+from rich.console import Console
+from know_lang_bot.utils.model_provider import create_pydantic_model
 
 LOG = FancyLogger(__name__)
+console = Console()
 
 class ChatStatus(str, Enum):
     """Enum for tracking chat progress status"""
@@ -128,7 +131,10 @@ Example Output: "Where is the configuration file or configuration settings store
     async def run(self, ctx: GraphRunContext[ChatGraphState, ChatGraphDeps]) -> RetrieveContextNode:
         # Create an agent for question polishing
         polish_agent = Agent(
-            f"{ctx.deps.config.llm.model_provider}:{ctx.deps.config.llm.model_name}",
+            create_pydantic_model(
+                model_provider=ctx.deps.config.llm.model_provider,
+                model_name=ctx.deps.config.llm.model_name
+            ),
             system_prompt=self.system_prompt
         )
         prompt = f"""Original question: "{ctx.state.original_question}"
@@ -203,7 +209,10 @@ Remember: Your primary goal is answering the user's specific question, not expla
 
     async def run(self, ctx: GraphRunContext[ChatGraphState, ChatGraphDeps]) -> End[ChatResult]:
         answer_agent = Agent(
-            f"{ctx.deps.config.llm.model_provider}:{ctx.deps.config.llm.model_name}",
+            create_pydantic_model(
+                model_provider=ctx.deps.config.llm.model_provider,
+                model_name=ctx.deps.config.llm.model_name
+            ),
             system_prompt=self.system_prompt
         )
         
@@ -267,6 +276,8 @@ async def process_chat(
         )
     except Exception as e:
         LOG.error(f"Error processing chat in graph: {e}")
+        console.print_exception()
+        
         result = ChatResult(
             answer="I encountered an error processing your question. Please try again."
         )

src/know_lang_bot/models/huggingface.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import AsyncIterator, Optional
+from contextlib import asynccontextmanager
+
+from pydantic_ai.models import (
+    AgentModel, 
+    ModelMessage, 
+    ModelResponse, 
+    StreamedResponse, 
+    Usage, 
+    Model,
+    check_allow_model_requests
+)
+from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.messages import (
+    ModelResponsePart, TextPart,  SystemPromptPart,
+    UserPromptPart, ToolReturnPart, ModelResponseStreamEvent
+)
+from pydantic_ai.settings import ModelSettings
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import torch
+
+
+@dataclass(init=False)
+class HuggingFaceModel(Model):
+    """A model that uses HuggingFace models locally.
+    
+    For MVP, this implements basic functionality without streaming or advanced features.
+    """
+    
+    model_name: str
+    model: AutoModelForCausalLM = field(repr=False)
+    tokenizer: AutoTokenizer = field(repr=False)
+    device: str = field(default="cuda" if torch.cuda.is_available() else "cpu")
+    
+    def __init__(
+        self,
+        model_name: str,
+        *,
+        device: Optional[str] = None,
+        max_new_tokens: int = 512,
+    ):
+        """Initialize a HuggingFace model.
+        
+        Args:
+            model_name: Name of the model on HuggingFace Hub
+            device: Device to run model on ('cuda' or 'cpu')
+            max_new_tokens: Maximum number of tokens to generate
+        """
+        self.model_name = model_name
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.max_new_tokens = max_new_tokens
+        
+        # Load model and tokenizer
+        self.model = AutoModelForCausalLM.from_pretrained(model_name)
+        self.model.to(self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        if not self.tokenizer.pad_token:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+    async def agent_model(
+        self,
+        *,
+        function_tools: list[ToolDefinition],
+        allow_text_result: bool,
+        result_tools: list[ToolDefinition],
+    ) -> AgentModel:
+        """Create an agent model for each step of an agent run."""
+        check_allow_model_requests()
+        
+        return HuggingFaceAgentModel(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            model_name=self.model_name,
+            device=self.device,
+            max_new_tokens=self.max_new_tokens,
+            tools=function_tools if function_tools else None
+        )
+
+    def name(self) -> str:
+        return f"huggingface:{self.model_name}"
+
+@dataclass
+class HuggingFaceAgentModel(AgentModel):
+    """Implementation of AgentModel for HuggingFace models."""
+    
+    model: AutoModelForCausalLM
+    tokenizer: AutoTokenizer
+    model_name: str
+    device: str
+    max_new_tokens: int
+    tools: Optional[list[ToolDefinition]] = None
+    
+    def _format_messages(self, messages: list[ModelMessage]) -> str:
+        """Format messages into a prompt the model can understand."""
+        formatted = []
+        for message in messages:
+            for part in message.parts:
+                if isinstance(part, SystemPromptPart):
+                    formatted.append(f"<|system|>{part.content}")
+                elif isinstance(part, UserPromptPart):
+                    formatted.append(f"<|user|>{part.content}")
+                else:
+                    # For MVP, we'll just pass through other message types
+                    formatted.append(str(part.content))
+        formatted.append("<|assistant|>")
+        return "\n".join(formatted)
+
+    async def request(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> tuple[ModelResponse, Usage]:
+        """Make a request to the model."""
+        prompt = self._format_messages(messages)
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=self.max_new_tokens,
+                pad_token_id=self.tokenizer.pad_token_id,
+            )
+        
+        response_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response_text = response_text[len(prompt):]  # Remove the prompt
+        
+        timestamp = datetime.now(timezone.utc)
+        response = ModelResponse(
+            parts=[TextPart(response_text)],
+            model_name=self.model_name,
+            timestamp=timestamp
+        )
+        
+        usage_stats = Usage(
+            requests=1,
+            request_tokens=inputs.input_ids.shape[1],
+            response_tokens=len(outputs[0]) - inputs.input_ids.shape[1],
+            total_tokens=len(outputs[0])
+        )
+        
+        return response, usage_stats
+
+    @asynccontextmanager
+    async def request_stream(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> AsyncIterator[StreamedResponse]:
+        """Make a streaming request to the model."""
+        prompt = self._format_messages(messages)
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        
+        streamer = TextIteratorStreamer(self.tokenizer, skip_special_tokens=True)
+        
+        generation_kwargs = dict(
+            input_ids=inputs.input_ids,
+            max_new_tokens=self.max_new_tokens,
+            pad_token_id=self.tokenizer.pad_token_id,
+            streamer=streamer,
+        )
+        
+        thread = torch.jit.fork(self.model.generate, **generation_kwargs)
+        
+        try:
+            yield HuggingFaceStreamedResponse(
+                _model_name=self.model_name,
+                _streamer=streamer,
+                _timestamp=datetime.now(timezone.utc),
+            )
+        finally:
+            torch.jit.wait(thread)
+
+@dataclass 
+class HuggingFaceStreamedResponse(StreamedResponse):
+    """Implementation of StreamedResponse for HuggingFace models."""
+    
+    _streamer: TextIteratorStreamer
+    _timestamp: datetime
+
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        """Stream tokens from the model."""
+        for new_text in self._streamer:
+            self._usage.response_tokens += len(new_text)
+            self._usage.total_tokens = self._usage.request_tokens + self._usage.response_tokens
+            
+            yield self._parts_manager.handle_text_delta(
+                vendor_part_id='content',
+                content=new_text
+            )
+
+    def timestamp(self) -> datetime:
+        return self._timestamp

src/know_lang_bot/summarizer/summarizer.py

@@ -10,6 +10,7 @@ from rich.progress import Progress
 from know_lang_bot.config import AppConfig
 from know_lang_bot.core.types import CodeChunk, ModelProvider
 from know_lang_bot.utils.fancy_log import FancyLogger
+from know_lang_bot.utils.model_provider import create_pydantic_model
 
 LOG = FancyLogger(__name__)
 
@@ -50,7 +51,10 @@ Provide a clean, concise and focused summary. Don't include unnecessary nor gene
 """
         
         self.agent = Agent(
-            f"{self.config.llm.model_provider}:{self.config.llm.model_name}",
+            create_pydantic_model(
+                model_provider=self.config.llm.model_provider,
+                model_name=self.config.llm.model_name
+            ),
             system_prompt=system_prompt,
             model_settings=self.config.llm.model_settings
         )

src/know_lang_bot/utils/model_provider.py

@@ -0,0 +1,17 @@
+from pydantic_ai.models import Model, KnownModelName
+from know_lang_bot.core.types import ModelProvider
+from know_lang_bot.models.huggingface import HuggingFaceModel
+from typing import get_args
+
+def create_pydantic_model(
+    model_provider: ModelProvider,
+    model_name: str,
+) -> Model | KnownModelName:
+    model_str = f"{model_provider}:{model_name}"
+
+    if model_str in get_args(KnownModelName):
+        return model_str
+    elif model_provider == ModelProvider.HUGGINGFACE:
+        return HuggingFaceModel(model_name=model_name)
+    else:
+        raise NotImplementedError(f"Model {model_provider}:{model_name} is not supported")