Chat, Embed, Rerank를 사용한 RAG 엔드투엔드 예시
Chat, Embed, Rerank를 사용한 RAG 엔드투엔드 예시
Chat, Embed, Rerank 엔드포인트(API v2)를 아우르는 Cohere의 검색 증강 생성(RAG) 능력 사용 가이드예요.
출처: 문서
본문
이 섹션은 기본 RAG 사용법을 확장해 다음을 포함하는 더 완전한 예시를 보여줍니다:
- 문서의 검색 및 재랭킹(Embed 및 Rerank 엔드포인트 경유).
- 챗봇용 RAG 구축(다중 턴 대화 포함).
설정(Setup)
먼저 Cohere 라이브러리를 임포트하고 클라이언트를 만듭니다.
Cohere 플랫폼
PYTHON
# ! pip install -U cohere
import cohere
import json
import numpy as np
co = cohere.ClientV2(
"COHERE_API_KEY"
) # Get your free API key here: https://dashboard.cohere.com/api-keys
프라이빗 배포(Private deployment)
PYTHON
# ! pip install -U cohere
import cohere
import json
import numpy as np
co = cohere.ClientV2(
api_key="", # Leave this blank
base_url="<YOUR_DEPLOYMENT_URL>",
)
1단계: 검색 쿼리 생성(Generating search queries)
다음으로, 사용자 질문에서 검색 쿼리를 생성하는 검색 쿼리 생성 도구를 만듭니다.
이 예시에서는 팀을 알게 되는 방법에 대해 묻는 사용자 쿼리를 전달해요.
PYTHON
message = "How to get to know my teammates"
# Define the query generation tool
query_gen_tool = [
{
"type": "function",
"function": {
"name": "internet_search",
"description": "Returns a list of relevant document snippets for a textual query retrieved from the internet",
"parameters": {
"type": "object",
"properties": {
"queries": {
"type": "array",
"items": {"type": "string"},
"description": "a list of queries to search the internet with.",
}
},
"required": ["queries"],
},
},
}
]
# Define a system message to optimize search query generation
instructions = "Write a search query that will find helpful information for answering the user's question accurately. If you need more than one search query, write a list of search queries. If you decide that a search is very unlikely to find information that would be useful in constructing a response to the user, you should instead directly answer."
# Generate search queries (if any)
search_queries = []
res = co.chat(
model="command-a-plus-05-2026",
messages=[
{"role": "system", "content": instructions},
{"role": "user", "content": message},
],
tools=query_gen_tool,
)
if res.message.tool_calls:
for tc in res.message.tool_calls:
queries = json.loads(tc.function.arguments)["queries"]
search_queries.extend(queries)
print(search_queries)
cURL
curl --request POST \
--url 'https://api.cohere.ai/v2/chat' \
--header 'accept: application/json' \
--header 'content-type: application/json' \
--header "Authorization: bearer ***" \
--data '{
"model": "command-a-plus-05-2026",
"messages": [
{
"role": "system",
"content": "Write a search query that will find helpful information for answering the user'\''s question accurately. If you need more than one search query, write a list of search queries. If you decide that a search is very unlikely to find information that would be useful in constructing a response to the user, you should instead directly answer."
},
{
"role": "user",
"content": "I'\''m joining a new team as a Principal Analyst. What are the best ways to quickly get to know my teammates?"
}
],
"tools": [
{
"type": "function",
"function": {
"name": "internet_search",
"description": "Returns a list of relevant document snippets for a textual query retrieved from the internet",
"parameters": {
"type": "object",
"properties": {
"queries": {
"type": "array",
"items": {
"type": "string"
},
"description": "a list of queries to search the internet with."
}
},
"required": ["queries"]
}
}
}
]
}'
예제 응답:
['how to get to know your teammates']
2단계: 관련 문서 가져오기(Fetching relevant documents)
Embed를 이용한 검색(Retrieval with Embed)
검색 쿼리가 주어졌으니, 대규모 문서 컬렉션에서 가장 관련 있는 문서를 검색할 방법이 필요해요.
여기서 Embed 엔드포인트를 통해 텍스트 임베딩을 활용할 수 있습니다.
Embed 엔드포인트는 의미론적 검색(semantic search)을 가능하게 해, 문서와 쿼리의 의미론적 의미를 비교할 수 있게 해줘요. 이는 키워드 일치를 찾는 데 뛰어나지만 텍스트 조각의 맥락이나 의미를 포착하는 데 어려움을 겪는, 더 전통적인 어휘 검색(lexical search) 접근 방식이 직면한 문제를 해결합니다.
Embed 엔드포인트는 텍스트를 입력으로 받아 임베딩을 출력으로 반환해요.
먼저 검색할 문서들을 임베딩해야 합니다. co.embed()를 사용해 Embed 엔드포인트를 호출하고 다음 인자를 전달해요:
model: 여기서는embed-v4.0을 선택합니다input_type: 검색 시 이들을 (쿼리가 아닌) 문서로 취급하도록search_document를 선택해요texts: 텍스트 목록(FAQ)embedding_types: 임베딩 유형으로float를 선택합니다.
PYTHON
# Define the documents
documents = [
{
"data": {
"text": "Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged."
}
},
{
"data": {
"text": "Finding Coffee Spots: For your caffeine fix, head to the break room's coffee machine or cross the street to the café for artisan coffee."
}
},
{
"data": {
"text": "Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!"
}
},
{
"data": {
"text": "Working Hours Flexibility: We prioritize work-life balance. While our core hours are 9 AM to 5 PM, we offer flexibility to adjust as needed."
}
},
{
"data": {
"text": "Side Projects Policy: We encourage you to pursue your passions. Just be mindful of any potential conflicts of interest with our business."
}
},
{
"data": {
"text": "Reimbursing Travel Expenses: Easily manage your travel expenses by submitting them through our finance tool. Approvals are prompt and straightforward."
}
},
{
"data": {
"text": "Working from Abroad: Working remotely from another country is possible. Simply coordinate with your manager and ensure your availability during core hours."
}
},
{
"data": {
"text": "Health and Wellness Benefits: We care about your well-being and offer gym memberships, on-site yoga classes, and comprehensive health insurance."
}
},
{
"data": {
"text": "Performance Reviews Frequency: We conduct informal check-ins every quarter and formal performance reviews twice a year."
}
},
{
"data": {
"text": "Proposing New Ideas: Innovation is welcomed! Share your brilliant ideas at our weekly team meetings or directly with your team lead."
}
},
]
# Embed the documents
doc_emb = co.embed(
model="embed-v4.0",
input_type="search_document",
texts=[doc["data"]["text"] for doc in documents],
embedding_types=["float"],
).embeddings.float
cURL
curl --request POST \
--url 'https://api.cohere.ai/v2/embed' \
--header 'accept: application/json' \
--header 'content-type: application/json' \
--header "Authorization: bearer ***" \
--data '{
"model": "embed-v4.0",
"input_type": "search_document",
"embedding_types": ["float"],
"texts": [
"Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.",
"Finding Coffee Spots: For your caffeine fix, head to the break room'\''s coffee machine or cross the street to the café for artisan coffee.",
"Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!",
"Working Hours Flexibility: We prioritize work-life balance. While our core hours are 9 AM to 5 PM, we offer flexibility to adjust as needed.",
"Side Projects Policy: We encourage you to pursue your passions. Just be mindful of any potential conflicts of interest with our business.",
"Reimbursing Travel Expenses: Easily manage your travel expenses by submitting them through our finance tool. Approvals are prompt and straightforward.",
"Working from Abroad: Working remotely from another country is possible. Simply coordinate with your manager and ensure your availability during core hours.",
"Health and Wellness Benefits: We care about your well-being and offer gym memberships, on-site yoga classes, and comprehensive health insurance.",
"Performance Reviews Frequency: We conduct informal check-ins every quarter and formal performance reviews twice a year.",
"Proposing New Ideas: Innovation is welcomed! Share your brilliant ideas at our weekly team meetings or directly with your team lead."
]
}'
Embed 엔드포인트 호출에서 search_query를 input_type으로 선택합니다. 이렇게 하면 모델이 검색 시 이 텍스트를 (문서가 아닌) 쿼리로 취급하게 됩니다.
PYTHON
# Embed the search query
query_emb = co.embed(
model="embed-v4.0",
input_type="search_query",
texts=search_queries,
embedding_types=["float"],
).embeddings.float
cURL
curl --request POST \
--url 'https://api.cohere.ai/v2/embed' \
--header 'accept: application/json' \
--header 'content-type: application/json' \
--header "Authorization: bearer ***" \
--data '{
"model": "embed-v4.0",
"input_type": "search_query",
"embedding_types": ["float"],
"texts": ["how to get to know your teammates"]
}'
이제 쿼리에 가장 관련 있는 문서를 검색하고 싶어요. 이를 위해 numpy 라이브러리를 사용해 내적(dot product) 접근 방식으로 각 쿼리-문서 쌍 사이의 유사성을 계산합니다.
각 쿼리-문서 쌍은 두 개가 얼마나 유사한지를 나타내는 점수를 반환해요. 그런 다음 이 점수를 내림차순으로 정렬하고 가장 유사한 상위 쌍을 선택하는데, 여기서는 5개를 선택합니다(임의의 선택이며 어떤 숫자든 선택할 수 있어요).
여기서는 유사성 점수와 함께 가장 관련 있는 문서를 보여줍니다.
PYTHON
# Compute dot product similarity and display results
n = 5
scores = np.dot(query_emb, np.transpose(doc_emb))[0]
max_idx = np.argsort(-scores)[:n]
retrieved_documents = [documents[item] for item in max_idx]
for rank, idx in enumerate(max_idx):
print(f"Rank: {rank+1}")
print(f"Score: {scores[idx]}")
print(f"Document: {retrieved_documents[rank]}\n")
Rank: 1
Score: 0.32653470360872655
Document: {'data': {'text': 'Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!'}}
Rank: 2
Score: 0.26851855352264786
Document: {'data': {'text': 'Proposing New Ideas: Innovation is welcomed! Share your brilliant ideas at our weekly team meetings or directly with your team lead.'}}
Rank: 3
Score: 0.2581341975304149
Document: {'data': {'text': 'Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.'}}
Rank: 4
Score: 0.18633336738178463
Document: {'data': {'text': "Finding Coffee Spots: For your caffeine fix, head to the break room's coffee machine or cross the street to the café for artisan coffee."}}
Rank: 5
Score: 0.13022396595682814
Document: {'data': {'text': 'Health and Wellness Benefits: We care about your well-being and offer gym memberships, on-site yoga classes, and comprehensive health insurance.'}}
참고(Note)
단순화를 위해 이 예시에서는 쿼리가 하나만 생성된다고 가정해요. 실용적인 구현에서는 여러 쿼리가 생성될 수 있습니다. 그런 시나리오에서는 각 쿼리에 대해 검색을 수행해야 해요.
Rerank로 재랭킹(Reranking with Rerank)
재랭킹은 의미론적 또는 어휘 검색의 결과를 더 끌어올릴 수 있어요. Rerank 엔드포인트는 검색 결과 목록을 받아 쿼리에 가장 관련 있는 문서 순으로 재랭킹합니다. 구현에는 단 한 줄의 코드만 필요해요.
co.rerank()를 사용해 엔드포인트를 호출하고 다음 인자를 전달합니다:
query: 사용자 쿼리documents: 의미론적 검색 결과에서 얻은 문서 목록top_n: 선택할 재랭킹된 상위 문서model: Rerank English 3을 선택합니다
결과를 보면, 쿼리가 팀을 알게 되는 것에 관한 것이므로 Slack 채널 가입에 관한 문서가 이전(3위)보다 상위(1위)로 올라간 것을 볼 수 있어요.
여기서 top_n을 2로 선택하는데, 이 문서들이 다음 응답 생성 단계로 전달될 것입니다.
PYTHON
# Rerank the documents
results = co.rerank(
model="rerank-v4.0-pro",
query=search_queries[0],
documents=[doc["data"]["text"] for doc in retrieved_documents],
top_n=2,
)
# Display the reranking results
for idx, result in enumerate(results.results):
print(f"Rank: {idx+1}")
print(f"Score: {result.relevance_score}")
print(f"Document: {retrieved_documents[result.index]}\n")
reranked_documents = [
retrieved_documents[result.index] for result in results.results
]
cURL
curl --request POST \
--url 'https://api.cohere.ai/v2/rerank' \
--header 'accept: application/json' \
--header 'content-type: application/json' \
--header "Authorization: bearer ***" \
--data '{
"model": "rerank-v4.0-pro",
"query": "how to get to know your teammates",
"top_n": 2,
"documents": [
"Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!",
"Proposing New Ideas: Innovation is welcomed! Share your brilliant ideas at our weekly team meetings or directly with your team lead.",
"Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.",
"Finding Coffee Spots: For your caffeine fix, head to the break room'\''s coffee machine or cross the street to the café for artisan coffee.",
"Health and Wellness Benefits: We care about your well-being and offer gym memberships, on-site yoga classes, and comprehensive health insurance."
]
}'
Rank: 1
Score: 0.07272241
Document: {'data': {'text': 'Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!'}}
Rank: 2
Score: 0.058674112
Document: {'data': {'text': 'Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.'}}
3단계: 응답 생성(Generating the response)
마지막으로, 검색된 문서를 전달하면서 Chat 엔드포인트를 호출해요. 이는 모델에게 RAG 모드로 실행하고 응답에 이 문서들을 사용하라고 알려줍니다.
그러면 쿼리와 검색된 문서를 기반으로 응답과 인용이 생성돼요.
PYTHON
messages = [{"role": "user", "content": message}]
# Generate the response
response = co.chat(
model="command-a-plus-05-2026",
messages=messages,
documents=reranked_documents,
)
# Display the response
print(response.message.content[0].text)
# Display the citations and source documents
if response.message.citations:
print("\nCITATIONS:")
for citation in response.message.citations:
print(citation, "\n")
cURL
curl --request POST \
--url 'https://api.cohere.ai/v2/chat' \
--header 'accept: application/json' \
--header 'content-type: application/json' \
--header "Authorization: bearer ***" \
--data '{
"model": "command-a-plus-05-2026",
"messages": [
{
"role": "user",
"content": "I'\''m joining a new team as a Principal Analyst. What are the best ways to quickly get to know my teammates?"
}
],
"documents": [
{
"data": {
"text": "Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!"
}
},
{
"data": {
"text": "Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged."
}
}
]
}'
To get to know your teammates, you can join relevant Slack channels to stay informed and engaged. You will receive an invite via email. You can also participate in team-building activities such as monthly outings and weekly game nights.
CITATIONS:
start=39 end=67 text='join relevant Slack channels' sources=[DocumentSource(type='document', id='doc:1', document={'id': 'doc:1', 'text': 'Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.'})] type='TEXT_CONTENT'
start=71 end=97 text='stay informed and engaged.' sources=[DocumentSource(type='document', id='doc:1', document={'id': 'doc:1', 'text': 'Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.'})] type='TEXT_CONTENT'
start=107 end=135 text='receive an invite via email.' sources=[DocumentSource(type='document', id='doc:1', document={'id': 'doc:1', 'text': 'Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.'})] type='TEXT_CONTENT'
start=164 end=188 text='team-building activities' sources=[DocumentSource(type='document', id='doc:0', document={'id': 'doc:0', 'text': 'Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!'})] type='TEXT_CONTENT'
start=197 end=236 text='monthly outings and weekly game nights.' sources=[DocumentSource(type='document', id='doc:0', document={'id': 'doc:0', 'text': 'Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!'})] type='TEXT_CONTENT'