app/core/generator.py

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()