Spaces:
Sleeping
Sleeping
File size: 1,761 Bytes
5dde853 |
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 |
import os
import chromadb
from llama_index.core import VectorStoreIndex
from llama_index.core.tools import QueryEngineTool
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
# from llama_index.llms.litellm import LiteLLM
from llama_index.core.agent.workflow import ReActAgent
def initialize_code_agent():
hf_token = os.environ.get('HF_TOKEN')
deepseek_token = os.environ.get('DEEPSEEK_TOKEN')
code_db = chromadb.PersistentClient(path="./code_db")
code_chroma_collection = code_db.get_or_create_collection('code')
code_vector_store = ChromaVectorStore(chroma_collection=code_chroma_collection)
embedding_model = HuggingFaceEmbedding(
model_name="BAAI/bge-small-en-v1.5",
device="cpu",
token=hf_token,
)
index = VectorStoreIndex.from_vector_store(code_vector_store, embed_model=embedding_model)
code_llm = HuggingFaceInferenceAPI(
model_name="deepseek-ai/deepseek-coder-1.3b-instruct",
api_key=deepseek_token,
token=hf_token,
)
code_query_engine = index.as_query_engine(
llm=code_llm,
similarity_top_k=3
)
code_query_engine_tool = QueryEngineTool.from_defaults(
query_engine=code_query_engine,
name="my_code_query_engine",
description="Code Query engine for the agent",
return_direct=False
)
return ReActAgent(
name="code_engine",
description="Query engine for the agent",
tools=[code_query_engine_tool],
system_prompt="You are a calculator assistant. Use your tools for any math operation.",
llm=code_llm
)
|