from dotenv import load_dotenv
from google.adk.agents import Agent
from sentence_transformers import SentenceTransformer
from .tools.tools import get_retrieval_tool
from .prompt import knowledge_agent_instruction
from pydantic import BaseModel, Field
load_dotenv()
# Sentence Transformer 모델 로드 (build_index.py와 동일한 모델 사용)
embedding_model = SentenceTransformer("jhgan/ko-sroberta-multitask")
Model = "gemini-2.5-flash-lite"
def get_embedding(text: str) -> list[float]:
"""
Sentence Transformer를 사용하여 텍스트의 임베딩을 생성합니다.
Args:
text: 임베딩을 생성할 텍스트
Returns:
텍스트의 임베딩 벡터
"""
embedding = embedding_model.encode(text, convert_to_tensor=False)
return embedding.tolist()
terms_rag_tool = get_retrieval_tool(
index_name="terms",
top_k=10,
)
code_package_rag_tool = get_retrieval_tool(
index_name="code_package",
top_k=10,
)
# 단어와 그 뜻의 dictionary로 저장
class KnowledgeOutput(BaseModel):
terms: dict[str, str] = Field(description="The terms and their explanations.")
root_agent = Agent(
name="knowledge_only",
model=Model,
description=
"This is agent that retrieves data from the faiss database"
"for understanding the information related to shipping and logistics.",
instruction= knowledge_agent_instruction,
tools=[terms_rag_tool, code_package_rag_tool],
output_key="knowledge_agent_output",
)
