calebdre

type QueryResult = {
	score: number,
	func: FunctionInfo
}
export async function runQuery(
	query: string,
	files: FileInfo[]
): Promise<QueryResult[]> {
	// convert input query to embedding
	const queryEmbedding = await createEmbedding(query);

calebdre

// math -> https://mathjs.org/index.html
export function cosineSimilarity(embedding1: number[], embedding2: number[]): number {
	const dotProduct = math.dot(embedding1, embedding2)
	const magnitude1 = math.norm(embedding1)
	const magnitude2 = math.norm(embedding2)
	return dotProduct / (magnitude1 * magnitude2)
}

calebdre

export const createEmbedding = async (
	model: string,
	text: string
) => {
	const resp = await ollama.embeddings({
		model: model,
		prompt: text
	})

	return resp.embedding

calebdre

const collectFileData = async (dir: string) => {
	// collect all the files in the directory
	const files = findTsFiles(dir)
	for (const file of files) {
		// extract functions from each file
		file.functions = extractFunctions(file.content)
		for (const func of file.functions) {
			// get the description for each function
			const description = await getFunctionDescription(file, func.functionName)
			if (description) {

calebdre

export const getFunctionDescription = async (file: FileInfo, functionName: string, model: string = "llama3"): Promise<string | null> => {
	const prompt = `given the following file ${file!.path}/${file!.filename}:
\`\`\`typescript
${file!.content}
\`\`\`

Provide a concise description for the function \`${functionName}\`. focusing on its specific role within the larger codebase.
The description should briefly explain:
- The purpose of the function
- The function's inputs, outputs, and any notable side effects or dependencies

calebdre

export function extractFunctions(fileContents: string): ExtractFunctionResponse[] {
	const functions: ExtractFunctionResponse[] = [];

	try {
		const ast = parser.parse(fileContents, {
			sourceType: 'module',
			plugins: ['typescript', 'jsx'],
			allowImportExportEverywhere: true,
			allowAwaitOutsideFunction: true,
			allowReturnOutsideFunction: true,

calebdre

type FileInfo {
  content: string
  path: string
  filename: string
}

function findTsFiles(directory: string, exclude_dirs: string[]): void {
 const fileList: FileInfo[] = []

 function traverseDirectory(currentDir: string) {

calebdre

const evaluateCodeSearch = async (
  files: FileInfo[],
  numSamples: number = 15,
  k: number = 5
) => {
  // create validation set
  const validationSet: ValidationQuery[] = await generateValidationQueries(files, numSamples)
  const resultSet: QueryResult[][] = []
  
  // run each validation query against the code search system

calebdre

export const calculateMAPAtK = (
  sampleQueries: SampleQuery[], 
  retrievedResults: QueryResult[][], 
  k: number
) => {
	const apScores = sampleQueries.map((query, index) => {
		return calculateAveragePrecision(query, retrievedResults[index], k);
	});

	return apScores.reduce((sum, score) => sum + score, 0) / apScores.length;

calebdre

export const calculateAveragePrecision = (
	sampleQuery: SampleQuery, 
	retrievedResults: QueryResult[], 
	k: number
) => {
	// get the function names from the validation set
	const relevantFunctions = sampleQuery.expectedFunctions.map((func) => func.functionName)

	// get the function names from the retrieved results
	const retrievedFunctionNames = retrievedResults.map((result) => result.funcName)