feat: solution for 'Повторный экзамен: Граф с рефлексией и доработкой'

src/llm_integration.py

@@ -0,0 +1,33 @@
+"""
+LLM integration module for LangChain with support for OpenAI and Ollama.
+Provides a reusable LLM client based on environment configuration.
+"""
+
+import os
+from typing import Union
+
+from langchain.llms import OpenAI, Ollama
+from langchain.chat_models import ChatOpenAI, ChatOllama
+
+# Environment variable to select provider: "openai" or "ollama"
+LLM_PROVIDER = os.getenv("LLM_PROVIDER", "openai").lower()
+
+
+def get_llm() -> Union[OpenAI, Ollama, ChatOpenAI, ChatOllama]:
+    """
+    Returns an LLM instance based on the configured provider.
+
+    For OpenAI, uses the default OpenAI LLM (text-davinci-003 or gpt-3.5-turbo).
+    For Ollama, uses the default Ollama LLM (e.g., llama2).
+
+    Raises:
+        ValueError: If an unsupported provider is specified.
+    """
+    if LLM_PROVIDER == "openai":
+        # Use ChatOpenAI for GPT-3.5-turbo by default
+        return ChatOpenAI(temperature=0.7)
+    elif LLM_PROVIDER == "ollama":
+        # Use ChatOllama for local models
+        return ChatOllama(model="llama2", temperature=0.7)
+    else:
+        raise ValueError(f"Unsupported LLM provider: {LLM_PROVIDER}")