25 lines
604 B
Python
25 lines
604 B
Python
from __future__ import annotations
|
|
|
|
from typing import List
|
|
|
|
from app.services.code_kb.adapter import CodeKnowledgeBaseAdapter
|
|
from app.services.code_kb.schema import CodeSearchHit
|
|
|
|
|
|
class CodeFunctionRetriever:
|
|
def __init__(self, adapter: CodeKnowledgeBaseAdapter) -> None:
|
|
self.adapter = adapter
|
|
|
|
def retrieve(
|
|
self,
|
|
query: str,
|
|
top_k: int = 8,
|
|
min_similarity: float = 0.0,
|
|
) -> List[CodeSearchHit]:
|
|
return self.adapter.search_functions(
|
|
query=query,
|
|
top_k=top_k,
|
|
min_similarity=min_similarity,
|
|
)
|
|
|