Merge pull request #8 from renansantosmendes/feat/add-chat-model

add chatpgl class and files

Author: renansantosmendes

.github/workflows/tests.yml

@@ -47,6 +47,8 @@ jobs:
         continue-on-error: true
 
       - name: Run tests
+        env:
+          GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
         run: uv run pytest tests/ -v --cov=pgl_utils --cov-report=xml --cov-report=term
 
       - name: Upload coverage reports

pgl_utils/genai/llm.py

@@ -3,8 +3,195 @@
 """
 
 
-def placeholder():
+# Standard library
+import os
+import time
+from typing import Any, Iterator, List, Optional
+
+# Third-party
+from groq import Groq
+from langchain_community.chat_message_histories import ChatMessageHistory
+from langchain_core.callbacks.manager import CallbackManagerForLLMRun
+from langchain_core.chat_history import BaseChatMessageHistory
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
+from langchain_core.outputs import ChatResult
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain_groq import ChatGroq
+from pydantic import ConfigDict
+
+
+_PGL_API_KEY = os.environ.get("GROQ_API_KEY", "")
+_FALLBACK_MODEL = "llama-3.3-70b-versatile"
+
+
+def _fetch_available_models() -> List:
+    client = Groq(api_key=_PGL_API_KEY)
+    models = client.models.list()
+    return sorted(models.data, key=lambda m: m.created, reverse=True)
+
+
+def _resolve_default_model() -> str:
+    try:
+        models = _fetch_available_models()
+        llama = next((m.id for m in models if "llama-3.3" in m.id), None)
+        return llama if llama else _FALLBACK_MODEL
+    except Exception:
+        return _FALLBACK_MODEL
+
+
+_DEFAULT_MODEL = _resolve_default_model()
+
+
+class ChatPGL(BaseChatModel):
     """
-    Placeholder function
+    ChatGroq wrapper pre-configured for PGL classes.
+    Students do not need to provide an API key.
+    Uses LangChain's ChatMessageHistory for conversation memory.
     """
-    return "GenAI LLM utilities"
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    model: str = _DEFAULT_MODEL
+    temperature: float = 0.7
+    max_tokens: int = 1024
+    stream_delay: float = 0.02
+    system_prompt: Optional[str] = None
+    session_id: str = "default"
+    _client: Any = None
+    _store: dict = {}
+    _chain: Any = None
+
+    def __init__(self, **kwargs):
+        if "model" in kwargs:
+            available = [m.id for m in _fetch_available_models()]
+            if kwargs["model"] not in available:
+                raise ValueError(
+                    f"Model '{kwargs['model']}' is not available.\n"
+                    f"Available models: {available}"
+                )
+        super().__init__(**kwargs)
+
+        self._store = {}
+        self._client = ChatGroq(
+            model=self.model,
+            temperature=self.temperature,
+            max_tokens=self.max_tokens,
+            api_key=_PGL_API_KEY,
+        )
+        self._chain = self._build_chain()
+
+    def _build_chain(self) -> RunnableWithMessageHistory:
+        system = self.system_prompt or "You are a helpful assistant."
+
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", system),
+            MessagesPlaceholder(variable_name="history"),
+            ("human", "{input}"),
+        ])
+
+        return RunnableWithMessageHistory(
+            prompt | self._client,
+            self._get_session_history,
+            input_messages_key="input",
+            history_messages_key="history",
+        )
+
+    def _get_session_history(self, session_id: str) -> BaseChatMessageHistory:
+        if session_id not in self._store:
+            self._store[session_id] = ChatMessageHistory()
+        return self._store[session_id]
+
+    def invoke(self, input: Any, session_id: Optional[str] = None, **kwargs) -> Any:
+        """
+        Runs inference using LangChain memory. Conversation history is
+        automatically managed per session_id.
+
+        Args:
+            input:      A string message.
+            session_id: Conversation session identifier. Defaults to self.session_id.
+        """
+        sid = session_id or self.session_id
+        return self._chain.invoke(
+            {"input": input},
+            config={"configurable": {"session_id": sid}},
+            **kwargs,
+        )
+
+    def clear_memory(self, session_id: Optional[str] = None):
+        """Clears the conversation memory for the given session."""
+        sid = session_id or self.session_id
+        if sid in self._store:
+            self._store[sid].clear()
+
+    def get_memory(self, session_id: Optional[str] = None) -> List[BaseMessage]:
+        """Returns the conversation history for the given session."""
+        sid = session_id or self.session_id
+        return self._get_session_history(sid).messages
+
+    @staticmethod
+    def list_models() -> List[str]:
+        """Fetches and returns the list of available models from the Groq API, sorted by most recent."""
+        return [m.id for m in _fetch_available_models()]
+
+    @property
+    def _llm_type(self) -> str:
+        return "chat-pgl"
+
+    @property
+    def _identifying_params(self) -> dict:
+        return {
+            "model": self.model,
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+            "stream_delay": self.stream_delay,
+        }
+
+    def _generate(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        return self._client._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
+
+    def _stream(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> Iterator[Any]:
+        yield from self._client._stream(messages, stop=stop, run_manager=run_manager, **kwargs)
+
+    def bind_tools(self, tools, **kwargs):
+        return self._client.bind_tools(tools, **kwargs)
+
+    def streamed_invoke(self, input: Any, session_id: Optional[str] = None, stream_delay: Optional[float] = None) -> str:
+        """
+        Runs inference with a typing effect, printing the response token by token.
+        Conversation history is automatically managed per session_id.
+        Returns the full response as a string when complete.
+
+        Args:
+            input:        A string message.
+            session_id:   Conversation session identifier. Defaults to self.session_id.
+            stream_delay: Seconds to wait between each chunk. Defaults to self.stream_delay.
+        """
+        sid = session_id or self.session_id
+        delay = stream_delay if stream_delay is not None else self.stream_delay
+        full_response = ""
+
+        for chunk in self._chain.stream(
+            {"input": input},
+            config={"configurable": {"session_id": sid}},
+        ):
+            token = chunk.content
+            print(token, end="", flush=True)
+            full_response += token
+            time.sleep(delay)
+
+        print()
+        return full_response

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pgl-utils"
-version = "0.1.8"
+version = "0.1.9"
 description = "Machine Learning, Deep Learning, and GenAI utilities for PUC and IBMEC post-graduation students"
 readme = "README.md"
 requires-python = ">=3.10"

requirements.txt

@@ -16,6 +16,12 @@ python-dotenv
 networkx>=2.8.0
 matplotlib>=3.5.0
 seaborn>=0.12.0
+
+# Generative AI (added based on llm.py imports)
+langchain-groq
+langchain-core
+langchain-community
+pydantic
 
 # Optional: Generative AI
 # openai>=0.27.0

setup.py

@@ -9,7 +9,7 @@
 
 setup(
     name="pgl-utils",
-    version="0.1.8",
+    version="0.1.9",
     author="Your Name",
     author_email="your.email@example.com",
     description="Machine Learning, Deep Learning, and GenAI utilities for Post-Graduation Lectures",

tests/conftest.py

@@ -1,2 +1,17 @@
 import matplotlib
 matplotlib.use('Agg')
+
+# groq.Groq must be mocked before pgl_utils.genai.llm is first imported,
+# because llm.py calls _resolve_default_model() at module level which hits the API.
+# conftest.py is loaded before any test file, so this pre-load is safe.
+from unittest.mock import MagicMock, patch as _patch
+
+_mock_model = MagicMock()
+_mock_model.id = "llama-3.3-70b-versatile"
+_mock_model.created = 1_000_000
+
+with _patch("groq.Groq") as _groq_mock:
+    _groq_mock.return_value.models.list.return_value.data = [_mock_model]
+    import pgl_utils.genai.llm  # noqa: E402
+
+del _mock_model, _groq_mock, _patch, MagicMock