Spaces:
Sleeping
Sleeping
File size: 1,817 Bytes
64d7fdf | 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 57 58 | from langchain_groq import ChatGroq
from langchain_core.messages import HumanMessage, SystemMessage
from app.config import config, settings
from app.utils.logger import logger
from typing import AsyncIterator
class LLMGenerator:
def __init__(self):
llm_config = config["models"]["llm"]
self.llm = ChatGroq(
model=llm_config["model_name"],
temperature=llm_config["temperature"],
max_tokens=llm_config["max_tokens"],
groq_api_key=settings.groq_api_key,
streaming=llm_config["streaming"]
)
logger.info(f"LLM initialized: {llm_config['model_name']}")
def generate(self, prompt: str, system_prompt: str = None) -> str:
messages = []
if system_prompt:
messages.append(SystemMessage(content=system_prompt))
messages.append(HumanMessage(content=prompt))
response = self.llm.invoke(messages)
return response.content
async def agenerate(self, prompt: str, system_prompt: str = None) -> str:
messages = []
if system_prompt:
messages.append(SystemMessage(content=system_prompt))
messages.append(HumanMessage(content=prompt))
response = await self.llm.ainvoke(messages)
return response.content
async def stream(self, prompt: str, system_prompt: str = None) -> AsyncIterator[str]:
messages = []
if system_prompt:
messages.append(SystemMessage(content=system_prompt))
messages.append(HumanMessage(content=prompt))
async for chunk in self.llm.astream(messages):
if chunk.content:
yield chunk.content
llm_generator = LLMGenerator()
|