luanpoppe commited on
Commit
20e3edd
·
1 Parent(s): baeaaa5

feat: refatorando código para também conseguir fazer requisções assíncronas com o claude

Browse files
_utils/gerar_relatorio_modelo_usuario/contextual_retriever.py CHANGED
@@ -1,7 +1,7 @@
1
  import os
2
  from langchain_openai import ChatOpenAI
3
  from typing import List, Dict, Tuple, Optional
4
- from anthropic import Anthropic
5
  import logging
6
  from langchain.schema import Document
7
  import asyncio
@@ -10,6 +10,7 @@ from typing import List
10
  from multiprocessing import Process, Barrier, Queue
11
  from dataclasses import dataclass
12
  from langchain_core.messages import HumanMessage
 
13
 
14
  from _utils.gerar_relatorio_modelo_usuario.llm_calls import claude_answer, gpt_answer
15
  from _utils.gerar_relatorio_modelo_usuario.prompts import contextual_prompt
@@ -21,12 +22,14 @@ from _utils.models.gerar_relatorio import (
21
 
22
  lista_contador = []
23
 
 
24
  class ContextualRetriever:
25
  def __init__(
26
  self, config: RetrievalConfig, claude_api_key: str, claude_context_model: str
27
  ):
28
  self.config = config
29
- self.claude_client = Anthropic(api_key=claude_api_key)
 
30
  self.logger = logging.getLogger(__name__)
31
  self.bm25 = None
32
  self.claude_context_model = claude_context_model
@@ -36,7 +39,10 @@ class ContextualRetriever:
36
  try:
37
  print("COMEÇOU A REQUISIÇÃO")
38
  prompt = contextual_prompt(full_text, chunk.content)
39
- # response = claude_answer(self.claude_client, self.claude_context_model, prompt)
 
 
 
40
  response = await gpt_answer(prompt)
41
  return response
42
  except Exception as e:
@@ -48,12 +54,15 @@ class ContextualRetriever:
48
  async def create_contextualized_chunk(self, chunk, full_text):
49
  lista_contador.append(0)
50
  print("contador: ", len(lista_contador))
51
- page_content = ""
52
- for i in range(
53
- max(0, chunk.page_number - 1),
54
- min(len(full_text), chunk.page_number + 2),
55
- ):
56
- page_content += full_text[i].page_content if full_text[i] else ""
 
 
 
57
 
58
  context = await self.llm_generate_context(page_content, chunk)
59
  return ContextualizedChunk(
@@ -70,6 +79,7 @@ class ContextualRetriever:
70
  ) -> List[ContextualizedChunk]:
71
  """Add context to all chunks"""
72
  contextualized_chunks = []
 
73
 
74
  async with asyncio.TaskGroup() as tg:
75
  tasks = [
 
1
  import os
2
  from langchain_openai import ChatOpenAI
3
  from typing import List, Dict, Tuple, Optional
4
+ from anthropic import Anthropic, AsyncAnthropic
5
  import logging
6
  from langchain.schema import Document
7
  import asyncio
 
10
  from multiprocessing import Process, Barrier, Queue
11
  from dataclasses import dataclass
12
  from langchain_core.messages import HumanMessage
13
+ from asgiref.sync import sync_to_async
14
 
15
  from _utils.gerar_relatorio_modelo_usuario.llm_calls import claude_answer, gpt_answer
16
  from _utils.gerar_relatorio_modelo_usuario.prompts import contextual_prompt
 
22
 
23
  lista_contador = []
24
 
25
+
26
  class ContextualRetriever:
27
  def __init__(
28
  self, config: RetrievalConfig, claude_api_key: str, claude_context_model: str
29
  ):
30
  self.config = config
31
+ # self.claude_client = Anthropic(api_key=claude_api_key)
32
+ self.claude_client = AsyncAnthropic(api_key=claude_api_key)
33
  self.logger = logging.getLogger(__name__)
34
  self.bm25 = None
35
  self.claude_context_model = claude_context_model
 
39
  try:
40
  print("COMEÇOU A REQUISIÇÃO")
41
  prompt = contextual_prompt(full_text, chunk.content)
42
+ # response = await claude_answer(
43
+ # self.claude_client, self.claude_context_model, prompt
44
+ # )
45
+
46
  response = await gpt_answer(prompt)
47
  return response
48
  except Exception as e:
 
54
  async def create_contextualized_chunk(self, chunk, full_text):
55
  lista_contador.append(0)
56
  print("contador: ", len(lista_contador))
57
+ # Código comentado abaixo é para ler as páginas ao redor da página atual do chunk
58
+ # page_content = ""
59
+ # for i in range(
60
+ # max(0, chunk.page_number - 1),
61
+ # min(len(full_text), chunk.page_number + 2),
62
+ # ):
63
+ # page_content += full_text[i].page_content if full_text[i] else ""
64
+ page_number = chunk.page_number - 1
65
+ page_content = full_text[page_number].page_content
66
 
67
  context = await self.llm_generate_context(page_content, chunk)
68
  return ContextualizedChunk(
 
79
  ) -> List[ContextualizedChunk]:
80
  """Add context to all chunks"""
81
  contextualized_chunks = []
82
+ lista_contador = []
83
 
84
  async with asyncio.TaskGroup() as tg:
85
  tasks = [
_utils/gerar_relatorio_modelo_usuario/llm_calls.py CHANGED
@@ -3,8 +3,11 @@ from langchain_core.messages import HumanMessage
3
  from langchain_openai import ChatOpenAI
4
 
5
 
6
- def claude_answer(claude_client, claude_context_model, prompt):
7
- response = claude_client.messages.create(
 
 
 
8
  model=claude_context_model,
9
  max_tokens=100,
10
  messages=[{"role": "user", "content": prompt}],
 
3
  from langchain_openai import ChatOpenAI
4
 
5
 
6
+ async def claude_answer(claude_client, claude_context_model, prompt):
7
+ print("\n")
8
+ print("Começou uma requisição pelo Claude")
9
+ print("\n")
10
+ response = await claude_client.messages.create(
11
  model=claude_context_model,
12
  max_tokens=100,
13
  messages=[{"role": "user", "content": prompt}],