Spaces:
Sleeping
Sleeping
File size: 2,364 Bytes
fb0495b d0eb2f2 fb0495b c9580eb fb0495b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
from retriever import Retriever
from reranker import Reranker
from anthropic import Anthropic
from typing import List
import os
retriever = Retriever()
reranker = Reranker()
class RAGAgent:
def __init__(
self,
retriever=retriever,
reranker=reranker,
anthropic_api_key: str = os.environ["anthropic_api_key"],
model: str = "claude-3-5-sonnet-20241022",
max_tokens: int = 1024,
temperature: float = 0.0,
):
self.retriever = retriever
self.reranker = reranker
self.client = Anthropic(api_key=anthropic_api_key)
self.model = model
self.max_tokens = max_tokens
self.temperature = temperature
def get_context(self, query: str) -> List[str]:
# Get initial candidates from retriever
retrieved_docs = self.retriever.search_similar(query)
# Rerank the candidates
context = self.reranker.rerank(query, retrieved_docs)
return context
def generate_prompt(self, context: List[str]) -> str:
context = "\n".join(context)
prompt = f"""
"讗转讛 专讜驻讗 砖讬谞讬讬诐, 讚讜讘专 注讘专讬转 讘诇讘讚. 拽讜专讗讬诐 诇讱 '专讜驻讗 讛砖讬谞讬讬诐 讛注讘专讬 讛讗诇拽讟专讜谞讬 讛专讗砖讜谉'. 注谞讛 诇诪讟讜驻诇 注诇 讛砖讗诇讛 砖诇讜 注诇 住诪讱 讛拽讜谞讟拽住 讛讘讗: {context}. 讛讜住祝 讻诪讛 砖讬讜转专 驻专讟讬诐, 讜讚讗讙 砖讛转讞讘讬专 讬讛讬讛 转拽讬谉 讜讬驻讛. 转注爪讜专 讻砖讗转讛 诪专讙讬砖 砖诪讬爪讬转 讗转 注爪诪讱. 讗诇 转诪爪讬讗 讚讘专讬诐. 讜讗诇 转注谞讛 讘砖驻讜转 砖讛谉 诇讗 注讘专讬转.
"""
return prompt
def get_response(self, question: str) -> str:
# Get relevant context
context = self.get_context(question)
# Generate prompt with context
prompt = self.generate_prompt(context)
# Get response from Claude
response = self.client.messages.create(
model=self.model,
max_tokens=self.max_tokens,
temperature=self.temperature,
messages=[
{"role": "assistant", "content": prompt},
{"role": "user", "content": f"{question}"}
]
)
return response.content[0].text
|