File size: 1,943 Bytes
575fee0 4bfba1f 575fee0 4bfba1f 575fee0 4bfba1f 575fee0 4bfba1f 575fee0 4bfba1f 575fee0 4bfba1f 575fee0 4bfba1f 575fee0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
from typing import Optional, Dict, Any, Union, Generator
from huggingface_hub import InferenceClient
from openai import OpenAI
from prompt_template import PromptTemplate
class AIAssistant:
def __init__(self, client: Union[OpenAI, InferenceClient], model: str):
self.client = client
self.model = model
def generate_response(
self,
prompt_template: PromptTemplate,
messages: list[Dict[str, str]],
generation_params: Optional[Dict] = None,
stream: bool = True,
) -> Generator[str, None, None]:
"""
Generate LLM response using the provided template and parameters.
Args:
prompt_template: PromptTemplate object containing template and parameters
messages: List of message dictionaries with role and content
generation_params: Optional generation parameters (overrides template parameters)
stream: Whether to stream the response
Yields:
Streamed response text
"""
params = generation_params or prompt_template.parameters
# Ensure messages are in correct format
formatted_messages = [
{"role": msg["role"], "content": str(msg["content"])}
for msg in messages
]
try:
completion = self.client.chat.completions.create(
model=self.model,
messages=formatted_messages,
stream=stream,
**params
)
if stream:
response = ""
for chunk in completion:
if chunk.choices[0].delta.content is not None:
response += chunk.choices[0].delta.content
yield response
else:
return completion.choices[0].message.content
except Exception as e:
yield f"Error generating response: {str(e)}" |