refactoring

analysis.py

@@ -0,0 +1,203 @@
+from llama_index.core.workflow import Workflow, Context, StartEvent, StopEvent, step
+import json
+from prompts import PRECEDENT_ANALYSIS_TEMPLATE
+from enum import Enum
+from anthropic import Anthropic
+from llama_index.llms.openai import OpenAI
+from llama_index.core.llms import ChatMessage
+from config import embed_model, Settings, openai_api_key, anthropic_api_key
+
+
+class ModelProvider(str, Enum):
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+
+
+class ModelName(str, Enum):
+    # OpenAI models
+    GPT4o = "gpt-4o"
+    GPT4o_MINI = "gpt-4o-mini"
+    # Anthropic models
+    CLAUDE3_5_SONNET = "claude-3-5-sonnet-latest"
+    CLAUDE3_5_HAIKU = "claude-3-5-haiku-latest"
+
+
+class LLMAnalyzer:
+    def __init__(self, provider: ModelProvider, model_name: ModelName):
+        self.provider = provider
+        self.model_name = model_name
+
+        if provider == ModelProvider.OPENAI:
+            self.client = OpenAI(model=model_name)
+        elif provider == ModelProvider.ANTHROPIC:
+            # Додаємо API ключ при ініціалізації
+            self.client = Anthropic(api_key=anthropic_api_key)
+        else:
+            raise ValueError(f"Unsupported provider: {provider}")
+
+    async def analyze(self, prompt: str, response_schema: dict) -> str:
+        if self.provider == ModelProvider.OPENAI:
+            return await self._analyze_with_openai(prompt, response_schema)
+        else:
+            return await self._analyze_with_anthropic(prompt, response_schema)
+
+    async def _analyze_with_openai(self, prompt: str, response_schema: dict) -> str:
+        messages = [
+            ChatMessage(role="system",
+                        content="Ти - кваліфікований юрист-аналітик, експерт з правових позицій Верховного Суду."),
+            ChatMessage(role="user", content=prompt)
+        ]
+
+        # Правильний формат для response_format
+        response_format = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": "relevant_positions_schema",  # Додаємо обов'язкове поле name
+                "schema": response_schema
+            }
+        }
+
+        response = self.client.chat(
+            messages=messages,
+            response_format=response_format,
+            temperature=0,
+            max_tokens=4096
+        )
+        return response.message.content
+
+    async def _analyze_with_anthropic(self, prompt: str, response_schema: dict) -> str:
+        response = self.client.messages.create(  # Прибрали await
+            model=self.model_name,
+            max_tokens=4096,
+            messages=[
+                {
+                    "role": "assistant",
+                    "content": "Ти - кваліфікований юрист-аналітик, експерт з правових позицій Верховного Суду."
+                },
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ]
+        )
+        return response.content[0].text
+
+
+class PrecedentAnalysisWorkflow(Workflow):
+    def __init__(self, provider: ModelProvider = ModelProvider.OPENAI,
+                 model_name: ModelName = ModelName.GPT4o_MINI):
+        super().__init__()
+        self.analyzer = LLMAnalyzer(provider, model_name)
+
+    @step
+    async def analyze(self, ctx: Context, ev: StartEvent) -> StopEvent:
+        try:
+            # Отримуємо параметри з події з дефолтними значеннями
+            query = ev.get("query", "")
+            question = ev.get("question", "")
+            nodes = ev.get("nodes", [])
+
+            # Перевірка на пусті значення
+            if not query:
+                return StopEvent(result="Помилка: Не надано текст нового рішення (query)")
+
+            if not nodes:
+                return StopEvent(result="Помилка: Не надано правові позиції для аналізу (nodes)")
+
+            # Підготовка контексту
+            context_parts = []
+            for i, node in enumerate(nodes, 1):
+                node_text = node.node.text if hasattr(node, 'node') else node.text
+                metadata = node.node.metadata if hasattr(node, 'node') else node.metadata
+                lp_id = metadata.get('lp_id', f'unknown_{i}')
+                context_parts.append(f"Source {i} (ID: {lp_id}):\n{node_text}")
+
+            context_str = "\n\n".join(context_parts)
+
+            # Схема відповіді
+            response_schema = {
+                "type": "object",
+                "properties": {
+                    "relevant_positions": {
+                        "type": "array",
+                        "items": {
+                            "type": "object",
+                            "properties": {
+                                "lp_id": {"type": "string"},
+                                "source_index": {"type": "string"},
+                                "description": {"type": "string"}
+                            },
+                            "required": ["lp_id", "source_index", "description"]
+                        }
+                    }
+                },
+                "required": ["relevant_positions"]
+            }
+
+            # Формування промпту
+            prompt = PRECEDENT_ANALYSIS_TEMPLATE.format(
+                query=query,
+                question=question if question else "Загальний аналіз релевантності",
+                context_str=context_str
+            )
+
+            # Отримання відповіді від моделі
+            response_content = await self.analyzer.analyze(prompt, response_schema)
+
+            try:
+                parsed_response = json.loads(response_content)
+                if "relevant_positions" in parsed_response:
+                    response_lines = []
+                    for position in parsed_response["relevant_positions"]:
+                        position_text = (
+                            f"* [{position['source_index']}] {position['description']}  "
+                        )
+                        response_lines.append(position_text)
+
+                    response_text = "\n".join(response_lines)
+                    return StopEvent(result=response_text)
+                else:
+                    return StopEvent(result="Не знайдено релевантних правових позицій")
+
+            except json.JSONDecodeError:
+                return StopEvent(result="Помилка обробки відповіді від AI")
+
+        except Exception as e:
+            return StopEvent(result=f"Error during analysis: {str(e)}")
+
+        # Формування промпту та отримання відповіді
+        prompt = PRECEDENT_ANALYSIS_TEMPLATE.format(
+            query=query,
+            question=question if question else "Загальний аналіз релевантності",
+            context_str=context_str
+        )
+
+        messages = [
+            ChatMessage(role="system", content="Ти - кваліфікований юрист-аналітик."),
+            ChatMessage(role="user", content=prompt)
+        ]
+
+        response = llm_analyse.chat(
+            messages=messages,
+            response_format=response_format
+        )
+
+        try:
+            parsed_response = json.loads(response.message.content)
+            if "relevant_positions" in parsed_response:
+                # Форматуємо результат
+                response_lines = []
+
+                for position in parsed_response["relevant_positions"]:
+                    position_text = (
+                        f"* [{position['source_index']}]: {position['description']} "
+                    )
+                    response_lines.append(position_text)
+
+                response_text = "\n".join(response_lines)
+                return StopEvent(result=response_text)
+            else:
+                return StopEvent(result="Помилка: відповідь не містить аналізу правових позицій")
+
+        except json.JSONDecodeError:
+            return StopEvent(result="Помилка обробки відповіді від AI")

config.py

@@ -0,0 +1,29 @@
+import os
+import google.generativeai as genai
+from llama_index.embeddings.openai import OpenAIEmbedding
+
+
+from llama_index.core import (
+    StorageContext,
+    ServiceContext,
+    VectorStoreIndex,
+    Settings,
+    load_index_from_storage
+)
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+aws_access_key_id = os.getenv("AWS_ACCESS_KEY_ID")
+aws_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY")
+openai_api_key = os.getenv("OPENAI_API_KEY")
+anthropic_api_key=os.getenv("ANTHROPIC_API_KEY")
+genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
+
+
+embed_model = OpenAIEmbedding(model_name="text-embedding-3-small")
+Settings.embed_model = embed_model
+Settings.context_window = 20000
+Settings.chunk_size = 2048
+Settings.similarity_top_k = 20

main.py

@@ -36,25 +36,28 @@ from llama_index.core.schema import NodeWithScore
 from llama_index.core.prompts import PromptTemplate
 from llama_index.core.response_synthesizers import ResponseMode, get_response_synthesizer
 
-from prompts import SYSTEM_PROMPT, LEGAL_POSITION_PROMPT, PRECEDENT_ANALYSIS_TEMPLATE
-
-
-from dotenv import load_dotenv
+from config import embed_model, Settings, openai_api_key, anthropic_api_key, aws_access_key_id, aws_secret_access_key
+from analysis import ModelProvider, ModelName, PrecedentAnalysisWorkflow
 
-load_dotenv()
-
-aws_access_key_id = os.getenv("AWS_ACCESS_KEY_ID")
-aws_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY")
-openai_api_key = os.getenv("OPENAI_API_KEY")
-anthropic_api_key=os.getenv("ANTHROPIC_API_KEY")
-genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
+from prompts import SYSTEM_PROMPT, LEGAL_POSITION_PROMPT, PRECEDENT_ANALYSIS_TEMPLATE
 
 
-embed_model = OpenAIEmbedding(model_name="text-embedding-3-small")
-Settings.embed_model = embed_model
-Settings.context_window = 20000
-Settings.chunk_size = 2048
-Settings.similarity_top_k = 20
+# from dotenv import load_dotenv
+#
+# load_dotenv()
+#
+# aws_access_key_id = os.getenv("AWS_ACCESS_KEY_ID")
+# aws_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY")
+# openai_api_key = os.getenv("OPENAI_API_KEY")
+# anthropic_api_key=os.getenv("ANTHROPIC_API_KEY")
+# genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
+#
+#
+# embed_model = OpenAIEmbedding(model_name="text-embedding-3-small")
+# Settings.embed_model = embed_model
+# Settings.context_window = 20000
+# Settings.chunk_size = 2048
+# Settings.similarity_top_k = 20
 
 
 # Параметри S3
@@ -117,202 +120,6 @@ state_lp_json = gr.State()
 state_nodes = gr.State()
 
 
-from enum import Enum
-
-
-class ModelProvider(str, Enum):
-    OPENAI = "openai"
-    ANTHROPIC = "anthropic"
-
-
-class ModelName(str, Enum):
-    # OpenAI models
-    GPT4o = "gpt-4o"
-    GPT4o_MINI = "gpt-4o-mini"
-    # Anthropic models
-    CLAUDE3_5_SONNET = "claude-3-5-sonnet-latest"
-    CLAUDE3_5_HAIKU = "claude-3-5-haiku-latest"
-
-
-class LLMAnalyzer:
-    def __init__(self, provider: ModelProvider, model_name: ModelName):
-        self.provider = provider
-        self.model_name = model_name
-
-        if provider == ModelProvider.OPENAI:
-            self.client = OpenAI(model=model_name)
-        elif provider == ModelProvider.ANTHROPIC:
-            # Додаємо API ключ при ініціалізації
-            self.client = Anthropic(api_key=anthropic_api_key)
-        else:
-            raise ValueError(f"Unsupported provider: {provider}")
-
-    async def analyze(self, prompt: str, response_schema: dict) -> str:
-        if self.provider == ModelProvider.OPENAI:
-            return await self._analyze_with_openai(prompt, response_schema)
-        else:
-            return await self._analyze_with_anthropic(prompt, response_schema)
-
-    async def _analyze_with_openai(self, prompt: str, response_schema: dict) -> str:
-        messages = [
-            ChatMessage(role="system",
-                        content="Ти - кваліфікований юрист-аналітик, експерт з правових позицій Верховного Суду."),
-            ChatMessage(role="user", content=prompt)
-        ]
-
-        # Правильний формат для response_format
-        response_format = {
-            "type": "json_schema",
-            "json_schema": {
-                "name": "relevant_positions_schema",  # Додаємо обов'язкове поле name
-                "schema": response_schema
-            }
-        }
-
-        response = self.client.chat(
-            messages=messages,
-            response_format=response_format,
-            temperature=0
-        )
-        return response.message.content
-
-    async def _analyze_with_anthropic(self, prompt: str, response_schema: dict) -> str:
-        response = self.client.messages.create(  # Прибрали await
-            model=self.model_name,
-            max_tokens=2000,
-            messages=[
-                {
-                    "role": "assistant",
-                    "content": "Ти - кваліфікований юрист-аналітик, експерт з правових позицій Верховного Суду."
-                },
-                {
-                    "role": "user",
-                    "content": prompt
-                }
-            ]
-        )
-        return response.content[0].text
-
-
-class PrecedentAnalysisWorkflow(Workflow):
-    def __init__(self, provider: ModelProvider = ModelProvider.OPENAI,
-                 model_name: ModelName = ModelName.GPT4o_MINI):
-        super().__init__()
-        self.analyzer = LLMAnalyzer(provider, model_name)
-
-    @step
-    async def analyze(self, ctx: Context, ev: StartEvent) -> StopEvent:
-        try:
-            # Отримуємо параметри з події з дефолтними значеннями
-            query = ev.get("query", "")
-            question = ev.get("question", "")
-            nodes = ev.get("nodes", [])
-
-            # Перевірка на пусті значення
-            if not query:
-                return StopEvent(result="Помилка: Не надано текст нового рішення (query)")
-
-            if not nodes:
-                return StopEvent(result="Помилка: Не надано правові позиції для аналізу (nodes)")
-
-            # Підготовка контексту
-            context_parts = []
-            for i, node in enumerate(nodes, 1):
-                node_text = node.node.text if hasattr(node, 'node') else node.text
-                metadata = node.node.metadata if hasattr(node, 'node') else node.metadata
-                lp_id = metadata.get('lp_id', f'unknown_{i}')
-                context_parts.append(f"Source {i} (ID: {lp_id}):\n{node_text}")
-
-            context_str = "\n\n".join(context_parts)
-
-            # Схема відповіді
-            response_schema = {
-                "type": "object",
-                "properties": {
-                    "relevant_positions": {
-                        "type": "array",
-                        "items": {
-                            "type": "object",
-                            "properties": {
-                                "lp_id": {"type": "string"},
-                                "source_index": {"type": "string"},
-                                "description": {"type": "string"}
-                            },
-                            "required": ["lp_id", "source_index", "description"]
-                        }
-                    }
-                },
-                "required": ["relevant_positions"]
-            }
-
-            # Формування промпту
-            prompt = PRECEDENT_ANALYSIS_TEMPLATE.format(
-                query=query,
-                question=question if question else "Загальний аналіз релевантності",
-                context_str=context_str
-            )
-
-            # Отримання відповіді від моделі
-            response_content = await self.analyzer.analyze(prompt, response_schema)
-
-            try:
-                parsed_response = json.loads(response_content)
-                if "relevant_positions" in parsed_response:
-                    response_lines = []
-                    for position in parsed_response["relevant_positions"]:
-                        position_text = (
-                            f"* [{position['source_index']}] {position['description']}  "
-                        )
-                        response_lines.append(position_text)
-
-                    response_text = "\n".join(response_lines)
-                    return StopEvent(result=response_text)
-                else:
-                    return StopEvent(result="Не знайдено релевантних правових позицій")
-
-            except json.JSONDecodeError:
-                return StopEvent(result="Помилка обробки відповіді від AI")
-
-        except Exception as e:
-            return StopEvent(result=f"Error during analysis: {str(e)}")
-
-        # Формування промпту та отримання відповіді
-        prompt = PRECEDENT_ANALYSIS_TEMPLATE.format(
-            query=query,
-            question=question if question else "Загальний аналіз релевантності",
-            context_str=context_str
-        )
-
-        messages = [
-            ChatMessage(role="system", content="Ти - кваліфікований юрист-аналітик."),
-            ChatMessage(role="user", content=prompt)
-        ]
-
-        response = llm_analyse.chat(
-            messages=messages,
-            response_format=response_format
-        )
-
-        try:
-            parsed_response = json.loads(response.message.content)
-            if "relevant_positions" in parsed_response:
-                # Форматуємо результат
-                response_lines = []
-
-                for position in parsed_response["relevant_positions"]:
-                    position_text = (
-                        f"* [{position['source_index']}]: {position['description']} "
-                    )
-                    response_lines.append(position_text)
-
-                response_text = "\n".join(response_lines)
-                return StopEvent(result=response_text)
-            else:
-                return StopEvent(result="Помилка: відповідь не містить аналізу правових позицій")
-
-        except json.JSONDecodeError:
-            return StopEvent(result="Помилка обробки відповіді від AI")
-
 
 def parse_doc_ids(doc_ids):
     if doc_ids is None: