@amplitude/ai 0.2.1 → 0.3.1

package/dist/mcp/validate-file.js

@@ -1,5 +1,451 @@
+import { require_acorn } from "../node_modules/.pnpm/acorn@8.16.0/node_modules/acorn/dist/acorn.js";
+import { require_lib } from "../node_modules/.pnpm/acorn-typescript@1.4.13_acorn@8.16.0/node_modules/acorn-typescript/lib/index.js";
+
 //#region src/mcp/validate-file.ts
-const llmPatterns = [
+const LLM_METHOD_CHAINS = [
+	{
+		chain: [
+			"chat",
+			"completions",
+			"create"
+		],
+		provider: "openai",
+		api: "chat.completions.create"
+	},
+	{
+		chain: [
+			"chat",
+			"completions",
+			"parse"
+		],
+		provider: "openai",
+		api: "chat.completions.parse"
+	},
+	{
+		chain: ["responses", "create"],
+		provider: "openai",
+		api: "responses.create"
+	},
+	{
+		chain: ["messages", "create"],
+		provider: "anthropic",
+		api: "messages.create"
+	},
+	{
+		chain: ["generateContent"],
+		provider: "gemini",
+		api: "generateContent",
+		singleMethodHints: [
+			"gemini",
+			"google",
+			"generative",
+			"genai"
+		]
+	},
+	{
+		chain: ["converse"],
+		provider: "bedrock",
+		api: "converse",
+		singleMethodHints: [
+			"bedrock",
+			"boto",
+			"runtime"
+		]
+	},
+	{
+		chain: ["chat", "complete"],
+		provider: "mistral",
+		api: "chat.complete"
+	},
+	{
+		chain: ["getChatCompletions"],
+		provider: "azure-openai",
+		api: "getChatCompletions",
+		singleMethodHints: ["azure", "openai"]
+	},
+	{
+		chain: ["chat"],
+		provider: "cohere",
+		api: "chat",
+		singleMethodHints: ["cohere", "CohereClient"]
+	},
+	{
+		chain: ["generate"],
+		provider: "cohere",
+		api: "generate",
+		singleMethodHints: ["cohere", "CohereClient"]
+	}
+];
+const VERCEL_AI_SDK_FUNCTIONS = new Set([
+	"streamText",
+	"generateText",
+	"streamObject",
+	"generateObject"
+]);
+const JS_KEYWORDS = new Set([
+	"if",
+	"else",
+	"for",
+	"while",
+	"do",
+	"switch",
+	"case",
+	"try",
+	"catch",
+	"finally",
+	"return",
+	"throw",
+	"new",
+	"delete",
+	"typeof",
+	"void",
+	"in",
+	"of",
+	"with",
+	"class",
+	"extends",
+	"super",
+	"import",
+	"export",
+	"default",
+	"break",
+	"continue",
+	"debugger",
+	"yield",
+	"await"
+]);
+const PROVIDER_CONSTRUCTOR_NAMES = new Set([
+	"OpenAI",
+	"Anthropic",
+	"Gemini",
+	"AzureOpenAI",
+	"Bedrock",
+	"Mistral",
+	"GoogleGenerativeAI",
+	"GoogleGenAI",
+	"CohereClient"
+]);
+function extractMethodChain(node) {
+	const parts = [];
+	let current = node;
+	while (current.type === "MemberExpression") {
+		const mem = current;
+		if (!mem.computed && mem.property.type === "Identifier") parts.push(mem.property.name);
+		current = mem.object;
+	}
+	parts.reverse();
+	return {
+		chain: parts,
+		receiver: current.type === "Identifier" ? current.name : null
+	};
+}
+function chainEndsWith(chain, pattern) {
+	if (chain.length < pattern.length) return false;
+	const offset = chain.length - pattern.length;
+	return pattern.every((p, i) => chain[offset + i] === p);
+}
+function hasSingleMethodHint(chain, receiver, hints) {
+	const tokens = [...chain.slice(0, -1)];
+	if (receiver) tokens.push(receiver);
+	return tokens.some((t) => hints.some((h) => t.toLowerCase().includes(h.toLowerCase())));
+}
+function getPropertyValue(obj, key) {
+	for (const prop of obj.properties) {
+		if (prop.type !== "Property") continue;
+		const p = prop;
+		if (p.key.type === "Identifier" && p.key.name === key) return p.value;
+		if (p.key.type === "Literal" && p.key.value === key) return p.value;
+	}
+	return null;
+}
+function hasKeywordArg(args, keyword) {
+	for (const arg of args) if (arg.type === "ObjectExpression") {
+		if (getPropertyValue(arg, keyword)) return true;
+	}
+	return false;
+}
+function extractCodeContext(sourceLines, lineNum, radius = 4) {
+	const idx = lineNum - 1;
+	const start = Math.max(0, idx - radius);
+	const end = Math.min(sourceLines.length - 1, idx + radius);
+	return sourceLines.slice(start, end + 1).map((l, i) => {
+		const ln = start + i + 1;
+		return `${ln === lineNum ? ">>>" : "   "} ${ln}: ${l}`;
+	}).join("\n");
+}
+function walkAST(node, callback, ancestors = []) {
+	callback(node, ancestors);
+	const next = [...ancestors, node];
+	for (const key of Object.keys(node)) {
+		if (key === "type" || key === "start" || key === "end" || key === "loc") continue;
+		const val = node[key];
+		if (val && typeof val === "object") {
+			if (Array.isArray(val)) {
+				for (const item of val) if (item && typeof item === "object" && "type" in item) walkAST(item, callback, next);
+			} else if ("type" in val) walkAST(val, callback, next);
+		}
+	}
+}
+function findContainingFunctionAST(ancestors) {
+	for (let i = ancestors.length - 1; i >= 0; i--) {
+		const node = ancestors[i];
+		if (!node) continue;
+		if (node.type === "FunctionDeclaration") {
+			const name = node.id?.name;
+			if (name && !JS_KEYWORDS.has(name)) return name;
+		}
+		if (node.type === "FunctionExpression") {
+			const fn = node;
+			if (fn.id?.name && !JS_KEYWORDS.has(fn.id.name)) return fn.id.name;
+		}
+		if (node.type === "VariableDeclarator") {
+			const decl = node;
+			if (decl.id.type === "Identifier" && decl.init && (decl.init.type === "ArrowFunctionExpression" || decl.init.type === "FunctionExpression")) {
+				const name = decl.id.name;
+				if (!JS_KEYWORDS.has(name)) return name;
+			}
+		}
+		if (node.type === "MethodDefinition" || node.type === "PropertyDefinition") {
+			const key = node.key;
+			if (key?.type === "Identifier") {
+				const name = key.name;
+				if (!JS_KEYWORDS.has(name)) return name;
+			}
+		}
+	}
+	return null;
+}
+function analyzeWithAST(source, sourceLines) {
+	let parse = null;
+	try {
+		const acorn = require_acorn();
+		const tsModule = require_lib();
+		const tsPlugin = typeof tsModule === "function" ? tsModule : tsModule.tsPlugin ?? tsModule.default;
+		if (typeof tsPlugin !== "function") return null;
+		parse = (src, opts) => acorn.Parser.extend(tsPlugin()).parse(src, opts);
+	} catch {
+		return null;
+	}
+	let tree;
+	try {
+		if (!parse) return null;
+		tree = parse(source, {
+			sourceType: "module",
+			ecmaVersion: "latest",
+			locations: true,
+			allowImportExportEverywhere: true,
+			allowReturnOutsideFunction: true
+		});
+	} catch {
+		return null;
+	}
+	const root = tree;
+	const hasPatch = /\bpatch\s*\(\s*\{/.test(source);
+	const hasSessionContext = /\.session\s*\(/.test(source) || /\bsession\.run\s*\(/.test(source) || /\.runAs\s*\(/.test(source) || /\bcreateAmplitudeAIMiddleware\s*\(/.test(source);
+	const hasAmplitudeImport = /from\s+['"]@amplitude\/ai['"]/.test(source) || /require\s*\(\s*['"]@amplitude\/ai['"]\s*\)/.test(source) || /\bAmplitudeAI\b/.test(source);
+	const wrappedClients = /* @__PURE__ */ new Set();
+	const localImportedBindings = /* @__PURE__ */ new Set();
+	const localImportedFunctions = /* @__PURE__ */ new Set();
+	const isLocalSpecifier = (spec) => spec.startsWith("./") || spec.startsWith("../") || spec.startsWith("@/") || spec.startsWith("~/");
+	walkAST(root, (node) => {
+		if (node.type === "ImportDeclaration") {
+			const decl = node;
+			const source$1 = decl.source;
+			if (!source$1 || source$1.type !== "Literal") return;
+			const specifier = source$1.value;
+			if (typeof specifier !== "string" || !isLocalSpecifier(specifier)) return;
+			const specifiers = decl.specifiers;
+			if (!specifiers) return;
+			for (const s of specifiers) {
+				const local = s.local;
+				if (local?.type === "Identifier") {
+					const name = local.name;
+					localImportedBindings.add(name);
+					localImportedFunctions.add(name);
+				}
+			}
+		}
+		if (node.type === "NewExpression") {
+			const ne = node;
+			const calleeName = ne.callee.type === "Identifier" ? ne.callee.name : null;
+			if (calleeName && PROVIDER_CONSTRUCTOR_NAMES.has(calleeName)) {
+				if (hasKeywordArg(ne.arguments, "amplitude")) {}
+			}
+		}
+	});
+	walkAST(root, (node) => {
+		if (node.type === "VariableDeclarator") {
+			const decl = node;
+			if (decl.id.type !== "Identifier" || !decl.init) return;
+			const varName = decl.id.name;
+			if (decl.init.type === "NewExpression") {
+				const ne = decl.init;
+				const calleeName = ne.callee.type === "Identifier" ? ne.callee.name : null;
+				if (calleeName && PROVIDER_CONSTRUCTOR_NAMES.has(calleeName) && hasKeywordArg(ne.arguments, "amplitude")) wrappedClients.add(varName);
+			}
+			if (decl.init.type === "CallExpression") {
+				const ce = decl.init;
+				if ((ce.callee.type === "Identifier" && ce.callee.name === "wrap" || ce.callee.type === "MemberExpression" && !ce.callee.computed && ce.callee.property.type === "Identifier" && ce.callee.property.name === "wrap") && hasAmplitudeImport) wrappedClients.add(varName);
+				if (hasAmplitudeImport && ce.callee.type === "Identifier") {
+					const fnName = ce.callee.name;
+					if (localImportedFunctions.has(fnName)) localImportedBindings.add(varName);
+				}
+			}
+		}
+	});
+	const callSites = [];
+	const assistantsApiRe = /\.beta\.(?:threads|assistants)\./;
+	const isReceiverInstrumented = (receiver) => {
+		if (hasPatch) return true;
+		if (receiver !== null && wrappedClients.has(receiver)) return true;
+		if (hasAmplitudeImport && hasSessionContext && receiver !== null && localImportedBindings.has(receiver)) return true;
+		return false;
+	};
+	walkAST(root, (node, ancestors) => {
+		if (node.type !== "CallExpression") return;
+		const ce = node;
+		const lineNum = node.loc?.start.line ?? 0;
+		if (lineNum === 0) return;
+		if (ce.callee.type === "Identifier") {
+			const fnName = ce.callee.name;
+			if (VERCEL_AI_SDK_FUNCTIONS.has(fnName)) callSites.push({
+				line: lineNum,
+				provider: "vercel-ai-sdk",
+				api: fnName,
+				instrumented: hasPatch,
+				containing_function: findContainingFunctionAST(ancestors),
+				code_context: extractCodeContext(sourceLines, lineNum)
+			});
+			return;
+		}
+		if (ce.callee.type === "MemberExpression") {
+			const mem = ce.callee;
+			if (!mem.computed && mem.property.type === "Identifier" && mem.property.name === "send" && ce.arguments.length > 0 && ce.arguments[0]?.type === "NewExpression") {
+				const newExpr = ce.arguments[0];
+				const cmdName = newExpr.callee.type === "Identifier" ? newExpr.callee.name : null;
+				if (cmdName === "InvokeModelCommand" || cmdName === "InvokeModelWithResponseStreamCommand") {
+					const receiver = mem.object.type === "Identifier" ? mem.object.name : null;
+					callSites.push({
+						line: lineNum,
+						provider: "bedrock",
+						api: cmdName === "InvokeModelWithResponseStreamCommand" ? "invokeModelWithResponseStream" : "invokeModel",
+						instrumented: isReceiverInstrumented(receiver),
+						containing_function: findContainingFunctionAST(ancestors),
+						code_context: extractCodeContext(sourceLines, lineNum)
+					});
+					return;
+				}
+				if (cmdName === "ConverseCommand" || cmdName === "ConverseStreamCommand") {
+					const receiver = mem.object.type === "Identifier" ? mem.object.name : null;
+					callSites.push({
+						line: lineNum,
+						provider: "bedrock",
+						api: cmdName === "ConverseStreamCommand" ? "converseStream" : "converse",
+						instrumented: isReceiverInstrumented(receiver),
+						containing_function: findContainingFunctionAST(ancestors),
+						code_context: extractCodeContext(sourceLines, lineNum)
+					});
+					return;
+				}
+			}
+		}
+		if (ce.callee.type === "MemberExpression") {
+			const { chain, receiver } = extractMethodChain(ce.callee);
+			for (const pattern of LLM_METHOD_CHAINS) {
+				if (!chainEndsWith(chain, pattern.chain)) continue;
+				if (pattern.chain.length === 1 && pattern.singleMethodHints) {
+					if (!hasSingleMethodHint(chain, receiver, pattern.singleMethodHints)) continue;
+				}
+				let effectiveProvider = pattern.provider;
+				let effectiveApi = pattern.api;
+				if (pattern.provider === "anthropic") {
+					const lineText = sourceLines[lineNum - 1] ?? "";
+					if (assistantsApiRe.test(lineText) || chain.some((p) => p === "beta" || p === "threads")) {
+						effectiveProvider = "openai-assistants";
+						effectiveApi = "beta.threads.messages.create";
+					}
+				}
+				callSites.push({
+					line: lineNum,
+					provider: effectiveProvider,
+					api: effectiveApi,
+					instrumented: isReceiverInstrumented(receiver),
+					containing_function: findContainingFunctionAST(ancestors),
+					code_context: extractCodeContext(sourceLines, lineNum)
+				});
+				break;
+			}
+		}
+	});
+	const toolDefinitions = [];
+	walkAST(root, (node) => {
+		if (node.type !== "ObjectExpression") return;
+		const obj = node;
+		const typeVal = getPropertyValue(obj, "type");
+		if (!typeVal || typeVal.type !== "Literal" || typeVal.value !== "function") return;
+		const fnVal = getPropertyValue(obj, "function");
+		if (!fnVal || fnVal.type !== "ObjectExpression") return;
+		const nameVal = getPropertyValue(fnVal, "name");
+		if (nameVal?.type === "Literal" && typeof nameVal.value === "string") {
+			const name = nameVal.value;
+			if (!toolDefinitions.includes(name)) toolDefinitions.push(name);
+		}
+	});
+	walkAST(root, (node) => {
+		if (node.type !== "ObjectExpression") return;
+		const obj = node;
+		const nameVal = getPropertyValue(obj, "name");
+		if (!nameVal || nameVal.type !== "Literal" || typeof nameVal.value !== "string") return;
+		const hasInputSchema = getPropertyValue(obj, "input_schema") !== null;
+		const hasParameters = getPropertyValue(obj, "parameters") !== null;
+		if (!hasInputSchema && !hasParameters) return;
+		const name = nameVal.value;
+		if (!toolDefinitions.includes(name)) toolDefinitions.push(name);
+	});
+	const functionDefinitions = [];
+	walkAST(root, (node) => {
+		if (node.type === "FunctionDeclaration") {
+			const name = node.id?.name;
+			if (name && !functionDefinitions.includes(name)) functionDefinitions.push(name);
+		}
+		if (node.type === "VariableDeclarator") {
+			const decl = node;
+			if (decl.id.type === "Identifier" && decl.init && (decl.init.type === "ArrowFunctionExpression" || decl.init.type === "FunctionExpression")) {
+				const name = decl.id.name;
+				if (!functionDefinitions.includes(name)) functionDefinitions.push(name);
+			}
+		}
+	});
+	const uninstrumentedSites = callSites.filter((s) => !s.instrumented);
+	let suggestions;
+	if (uninstrumentedSites.length > 0) suggestions = [
+		"Now: instrument provider calls with wrapper/swap import (for example, `new OpenAI({ amplitude: ai, apiKey })`).",
+		"Next: add session lineage with `const session = ai.agent(...).session(...)` and wrap calls in `session.run(...)`.",
+		"Why: session context unlocks session-level scoring, enrichments, and reliable product-to-AI funnels.",
+		"Content tiers: choose `contentMode` (`full`, `metadata_only`, or `customer_enriched`) and prefer `redactPii: true` when using `full`.",
+		"Fallback: use `patch({ amplitudeAI: ai })` for migration speed, then graduate to wrapper + session context."
+	];
+	else if (callSites.length > 0 && !hasSessionContext) suggestions = [
+		"Tracking is present, but session context is missing.",
+		"Now: add `const session = ai.agent(...).session(...)` and execute LLM calls inside `session.run(...)` (or use middleware).",
+		"Why: without session lineage you lose high-value outcomes like session enrichments, scoring, and dependable session funnels.",
+		"Content tiers: set `contentMode` intentionally and keep `redactPii: true` when using `full`."
+	];
+	else if (callSites.length > 0) suggestions = ["File appears instrumented with session lineage.", "For privacy-by-default, set `contentMode` intentionally and enable `redactPii: true` when using `full`."];
+	else suggestions = ["No supported LLM call sites detected in this file.", "If this file should emit AI telemetry, add wrapped provider calls and session context."];
+	return {
+		total_call_sites: callSites.length,
+		instrumented: callSites.length - uninstrumentedSites.length,
+		uninstrumented: uninstrumentedSites.length,
+		has_amplitude_import: hasAmplitudeImport,
+		has_session_context: hasSessionContext,
+		call_sites: callSites,
+		suggestions,
+		tool_definitions: toolDefinitions,
+		function_definitions: functionDefinitions
+	};
+}
+const regexLlmPatterns = [
 	{
 		pattern: /\.chat\.completions\.create\s*\(/g,
 		receiverRe: /(\w+)(?:\.\w+)*\.chat\.completions\.create\s*\(/,
@@ -41,11 +487,59 @@ const llmPatterns = [
 		receiverRe: /(\w+)(?:\.\w+)*\.send\s*\(/,
 		provider: "bedrock",
 		api: "converse"
+	},
+	{
+		pattern: /\.send\s*\(\s*new\s+InvokeModelWithResponseStreamCommand\s*\(/g,
+		receiverRe: /(\w+)(?:\.\w+)*\.send\s*\(/,
+		provider: "bedrock",
+		api: "invokeModelWithResponseStream"
+	},
+	{
+		pattern: /\.send\s*\(\s*new\s+ConverseStreamCommand\s*\(/g,
+		receiverRe: /(\w+)(?:\.\w+)*\.send\s*\(/,
+		provider: "bedrock",
+		api: "converseStream"
+	},
+	{
+		pattern: /\.chat\.complete\s*\(/g,
+		receiverRe: /(\w+)(?:\.\w+)*\.chat\.complete\s*\(/,
+		provider: "mistral",
+		api: "chat.complete"
+	},
+	{
+		pattern: /\.getChatCompletions\s*\(/g,
+		receiverRe: /(\w+)(?:\.\w+)*\.getChatCompletions\s*\(/,
+		provider: "azure-openai",
+		api: "getChatCompletions"
+	},
+	{
+		pattern: /\bstreamText\s*\(/g,
+		receiverRe: null,
+		provider: "vercel-ai-sdk",
+		api: "streamText"
+	},
+	{
+		pattern: /\bgenerateText\s*\(/g,
+		receiverRe: null,
+		provider: "vercel-ai-sdk",
+		api: "generateText"
+	},
+	{
+		pattern: /\bstreamObject\s*\(/g,
+		receiverRe: null,
+		provider: "vercel-ai-sdk",
+		api: "streamObject"
+	},
+	{
+		pattern: /\bgenerateObject\s*\(/g,
+		receiverRe: null,
+		provider: "vercel-ai-sdk",
+		api: "generateObject"
 	}
 ];
-function findWrappedConstructors(source) {
+function findWrappedConstructorsRegex(source) {
 	const result = /* @__PURE__ */ new Set();
-	for (const m of source.matchAll(/(?:const|let|var)\s+(\w+)\s*=\s*new\s+(?:OpenAI|Anthropic|Gemini|AzureOpenAI|Bedrock|Mistral)\s*\(/g)) {
+	for (const m of source.matchAll(/(?:const|let|var)\s+(\w+)\s*=\s*new\s+(?:OpenAI|Anthropic|Gemini|AzureOpenAI|Bedrock|Mistral|GoogleGenerativeAI|GoogleGenAI|CohereClient)\s*\(/g)) {
 		const varName = m[1] ?? "";
 		if (!varName) continue;
 		let depth = 1;
@@ -62,33 +556,80 @@ function findWrappedConstructors(source) {
 	}
 	return result;
 }
-function analyzeFileInstrumentation(source) {
+function findContainingFunctionRegex(lines, lineIndex) {
+	for (let i = lineIndex; i >= 0; i--) {
+		const line = lines[i] ?? "";
+		const fnDeclMatch = line.match(/(?:export\s+)?(?:async\s+)?function\s+(\w+)\s*\(/);
+		if (fnDeclMatch) {
+			const name = fnDeclMatch[1] ?? "";
+			if (name && !JS_KEYWORDS.has(name)) return name;
+		}
+		const arrowMatch = line.match(/(?:export\s+)?(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?\(/);
+		if (arrowMatch) {
+			const name = arrowMatch[1] ?? "";
+			if (name && !JS_KEYWORDS.has(name)) return name;
+		}
+	}
+	return null;
+}
+function findToolDefinitionsRegex(source) {
+	const tools = [];
+	for (const m of source.matchAll(/function:\s*\{[^}]*name:\s*['"](\w+)['"]/g)) {
+		const name = m[1] ?? "";
+		if (name && !tools.includes(name)) tools.push(name);
+	}
+	for (const m of source.matchAll(/name:\s*['"](\w+)['"][^}]*input_schema\s*:/g)) {
+		const name = m[1] ?? "";
+		if (name && !tools.includes(name)) tools.push(name);
+	}
+	return tools;
+}
+function findFunctionDefinitionsRegex(source) {
+	const fns = [];
+	for (const m of source.matchAll(/(?:async\s+)?function\s+(\w+)|(?:export\s+)?(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?(?:\([^)]*\)|[^=])\s*=>/g)) {
+		const name = m[1] ?? m[2] ?? "";
+		if (name && !fns.includes(name)) fns.push(name);
+	}
+	return fns;
+}
+function analyzeWithRegex(source) {
 	const hasPatch = /\bpatch\s*\(\s*\{/.test(source);
 	const hasSessionContext = /\.session\s*\(/.test(source) || /\bsession\.run\s*\(/.test(source) || /\.runAs\s*\(/.test(source) || /\bcreateAmplitudeAIMiddleware\s*\(/.test(source);
 	const hasAmplitudeImport = /from\s+['"]@amplitude\/ai['"]/.test(source) || /require\s*\(\s*['"]@amplitude\/ai['"]\s*\)/.test(source) || /\bAmplitudeAI\b/.test(source);
-	const wrappedClients = findWrappedConstructors(source);
+	const wrappedClients = findWrappedConstructorsRegex(source);
 	if (hasAmplitudeImport) for (const m of source.matchAll(/(?:const|let|var)\s+(\w+)\s*=\s*(?:\w+\.)*wrap\s*\(/g)) {
 		const name = m[1] ?? "";
 		if (name) wrappedClients.add(name);
 	}
 	const callSites = [];
 	const lines = source.split("\n");
+	const assistantsApiRe = /\.beta\.(?:threads|assistants)\./;
 	for (let i = 0; i < lines.length; i++) {
 		const line = lines[i] ?? "";
-		for (const { pattern, receiverRe, provider, api } of llmPatterns) {
+		for (const { pattern, receiverRe, provider, api } of regexLlmPatterns) {
 			pattern.lastIndex = 0;
 			if (pattern.test(line)) {
-				const receiver = line.match(receiverRe)?.[1] ?? "";
+				let effectiveProvider = provider;
+				let effectiveApi = api;
+				if (provider === "anthropic" && assistantsApiRe.test(line)) {
+					effectiveProvider = "openai-assistants";
+					effectiveApi = "beta.threads.messages.create";
+				}
+				const receiver = (receiverRe ? line.match(receiverRe) : null)?.[1] ?? "";
 				const instrumented = hasPatch || wrappedClients.has(receiver);
 				callSites.push({
 					line: i + 1,
-					provider,
-					api,
-					instrumented
+					provider: effectiveProvider,
+					api: effectiveApi,
+					instrumented,
+					containing_function: findContainingFunctionRegex(lines, i),
+					code_context: extractCodeContext(lines, i + 1)
 				});
 			}
 		}
 	}
+	const toolDefs = findToolDefinitionsRegex(source);
+	const fnDefs = findFunctionDefinitionsRegex(source);
 	const uninstrumentedSites = callSites.filter((s) => !s.instrumented);
 	let suggestions;
 	if (uninstrumentedSites.length > 0) suggestions = [
@@ -113,9 +654,16 @@ function analyzeFileInstrumentation(source) {
 		has_amplitude_import: hasAmplitudeImport,
 		has_session_context: hasSessionContext,
 		call_sites: callSites,
-		suggestions
+		suggestions,
+		tool_definitions: toolDefs,
+		function_definitions: fnDefs
 	};
 }
+function analyzeFileInstrumentation(source) {
+	const astResult = analyzeWithAST(source, source.split("\n"));
+	if (astResult) return astResult;
+	return analyzeWithRegex(source);
+}
 
 //#endregion
 export { analyzeFileInstrumentation };