fast-agent-mcp 0.1.12__py3-none-any.whl → 0.2.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {fast_agent_mcp-0.1.12.dist-info → fast_agent_mcp-0.2.0.dist-info}/METADATA +3 -4
- fast_agent_mcp-0.2.0.dist-info/RECORD +123 -0
- mcp_agent/__init__.py +75 -0
- mcp_agent/agents/agent.py +61 -415
- mcp_agent/agents/base_agent.py +522 -0
- mcp_agent/agents/workflow/__init__.py +1 -0
- mcp_agent/agents/workflow/chain_agent.py +173 -0
- mcp_agent/agents/workflow/evaluator_optimizer.py +362 -0
- mcp_agent/agents/workflow/orchestrator_agent.py +591 -0
- mcp_agent/{workflows/orchestrator → agents/workflow}/orchestrator_models.py +11 -21
- mcp_agent/agents/workflow/parallel_agent.py +182 -0
- mcp_agent/agents/workflow/router_agent.py +307 -0
- mcp_agent/app.py +15 -19
- mcp_agent/cli/commands/bootstrap.py +19 -38
- mcp_agent/cli/commands/config.py +4 -4
- mcp_agent/cli/commands/setup.py +7 -14
- mcp_agent/cli/main.py +7 -10
- mcp_agent/cli/terminal.py +3 -3
- mcp_agent/config.py +25 -40
- mcp_agent/context.py +12 -21
- mcp_agent/context_dependent.py +3 -5
- mcp_agent/core/agent_types.py +10 -7
- mcp_agent/core/direct_agent_app.py +179 -0
- mcp_agent/core/direct_decorators.py +443 -0
- mcp_agent/core/direct_factory.py +476 -0
- mcp_agent/core/enhanced_prompt.py +23 -55
- mcp_agent/core/exceptions.py +8 -8
- mcp_agent/core/fastagent.py +145 -371
- mcp_agent/core/interactive_prompt.py +424 -0
- mcp_agent/core/mcp_content.py +17 -17
- mcp_agent/core/prompt.py +6 -9
- mcp_agent/core/request_params.py +6 -3
- mcp_agent/core/validation.py +92 -18
- mcp_agent/executor/decorator_registry.py +9 -17
- mcp_agent/executor/executor.py +8 -17
- mcp_agent/executor/task_registry.py +2 -4
- mcp_agent/executor/temporal.py +19 -41
- mcp_agent/executor/workflow.py +3 -5
- mcp_agent/executor/workflow_signal.py +15 -21
- mcp_agent/human_input/handler.py +4 -7
- mcp_agent/human_input/types.py +2 -3
- mcp_agent/llm/__init__.py +2 -0
- mcp_agent/llm/augmented_llm.py +450 -0
- mcp_agent/llm/augmented_llm_passthrough.py +162 -0
- mcp_agent/llm/augmented_llm_playback.py +83 -0
- mcp_agent/llm/memory.py +103 -0
- mcp_agent/{workflows/llm → llm}/model_factory.py +22 -16
- mcp_agent/{workflows/llm → llm}/prompt_utils.py +1 -3
- mcp_agent/llm/providers/__init__.py +8 -0
- mcp_agent/{workflows/llm → llm/providers}/anthropic_utils.py +8 -25
- mcp_agent/{workflows/llm → llm/providers}/augmented_llm_anthropic.py +56 -194
- mcp_agent/llm/providers/augmented_llm_deepseek.py +53 -0
- mcp_agent/{workflows/llm → llm/providers}/augmented_llm_openai.py +99 -190
- mcp_agent/{workflows/llm → llm}/providers/multipart_converter_anthropic.py +72 -71
- mcp_agent/{workflows/llm → llm}/providers/multipart_converter_openai.py +65 -71
- mcp_agent/{workflows/llm → llm}/providers/openai_multipart.py +16 -44
- mcp_agent/{workflows/llm → llm/providers}/openai_utils.py +4 -4
- mcp_agent/{workflows/llm → llm}/providers/sampling_converter_anthropic.py +9 -11
- mcp_agent/{workflows/llm → llm}/providers/sampling_converter_openai.py +8 -12
- mcp_agent/{workflows/llm → llm}/sampling_converter.py +3 -31
- mcp_agent/llm/sampling_format_converter.py +37 -0
- mcp_agent/logging/events.py +1 -5
- mcp_agent/logging/json_serializer.py +7 -6
- mcp_agent/logging/listeners.py +20 -23
- mcp_agent/logging/logger.py +17 -19
- mcp_agent/logging/rich_progress.py +10 -8
- mcp_agent/logging/tracing.py +4 -6
- mcp_agent/logging/transport.py +22 -22
- mcp_agent/mcp/gen_client.py +1 -3
- mcp_agent/mcp/interfaces.py +117 -110
- mcp_agent/mcp/logger_textio.py +97 -0
- mcp_agent/mcp/mcp_agent_client_session.py +7 -7
- mcp_agent/mcp/mcp_agent_server.py +8 -8
- mcp_agent/mcp/mcp_aggregator.py +102 -143
- mcp_agent/mcp/mcp_connection_manager.py +20 -27
- mcp_agent/mcp/prompt_message_multipart.py +68 -16
- mcp_agent/mcp/prompt_render.py +77 -0
- mcp_agent/mcp/prompt_serialization.py +30 -48
- mcp_agent/mcp/prompts/prompt_constants.py +18 -0
- mcp_agent/mcp/prompts/prompt_helpers.py +327 -0
- mcp_agent/mcp/prompts/prompt_load.py +109 -0
- mcp_agent/mcp/prompts/prompt_server.py +155 -195
- mcp_agent/mcp/prompts/prompt_template.py +35 -66
- mcp_agent/mcp/resource_utils.py +7 -14
- mcp_agent/mcp/sampling.py +17 -17
- mcp_agent/mcp_server/agent_server.py +13 -17
- mcp_agent/mcp_server_registry.py +13 -22
- mcp_agent/resources/examples/{workflows → in_dev}/agent_build.py +3 -2
- mcp_agent/resources/examples/in_dev/slides.py +110 -0
- mcp_agent/resources/examples/internal/agent.py +6 -3
- mcp_agent/resources/examples/internal/fastagent.config.yaml +8 -2
- mcp_agent/resources/examples/internal/job.py +2 -1
- mcp_agent/resources/examples/internal/prompt_category.py +1 -1
- mcp_agent/resources/examples/internal/prompt_sizing.py +3 -5
- mcp_agent/resources/examples/internal/sizer.py +2 -1
- mcp_agent/resources/examples/internal/social.py +2 -1
- mcp_agent/resources/examples/prompting/agent.py +2 -1
- mcp_agent/resources/examples/prompting/image_server.py +4 -8
- mcp_agent/resources/examples/prompting/work_with_image.py +19 -0
- mcp_agent/ui/console_display.py +16 -20
- fast_agent_mcp-0.1.12.dist-info/RECORD +0 -161
- mcp_agent/core/agent_app.py +0 -646
- mcp_agent/core/agent_utils.py +0 -71
- mcp_agent/core/decorators.py +0 -455
- mcp_agent/core/factory.py +0 -463
- mcp_agent/core/proxies.py +0 -269
- mcp_agent/core/types.py +0 -24
- mcp_agent/eval/__init__.py +0 -0
- mcp_agent/mcp/stdio.py +0 -111
- mcp_agent/resources/examples/data-analysis/analysis-campaign.py +0 -188
- mcp_agent/resources/examples/data-analysis/analysis.py +0 -65
- mcp_agent/resources/examples/data-analysis/fastagent.config.yaml +0 -41
- mcp_agent/resources/examples/data-analysis/mount-point/WA_Fn-UseC_-HR-Employee-Attrition.csv +0 -1471
- mcp_agent/resources/examples/mcp_researcher/researcher-eval.py +0 -53
- mcp_agent/resources/examples/researcher/fastagent.config.yaml +0 -66
- mcp_agent/resources/examples/researcher/researcher-eval.py +0 -53
- mcp_agent/resources/examples/researcher/researcher-imp.py +0 -190
- mcp_agent/resources/examples/researcher/researcher.py +0 -38
- mcp_agent/resources/examples/workflows/chaining.py +0 -44
- mcp_agent/resources/examples/workflows/evaluator.py +0 -78
- mcp_agent/resources/examples/workflows/fastagent.config.yaml +0 -24
- mcp_agent/resources/examples/workflows/human_input.py +0 -25
- mcp_agent/resources/examples/workflows/orchestrator.py +0 -73
- mcp_agent/resources/examples/workflows/parallel.py +0 -78
- mcp_agent/resources/examples/workflows/router.py +0 -53
- mcp_agent/resources/examples/workflows/sse.py +0 -23
- mcp_agent/telemetry/__init__.py +0 -0
- mcp_agent/telemetry/usage_tracking.py +0 -18
- mcp_agent/workflows/__init__.py +0 -0
- mcp_agent/workflows/embedding/__init__.py +0 -0
- mcp_agent/workflows/embedding/embedding_base.py +0 -61
- mcp_agent/workflows/embedding/embedding_cohere.py +0 -49
- mcp_agent/workflows/embedding/embedding_openai.py +0 -46
- mcp_agent/workflows/evaluator_optimizer/__init__.py +0 -0
- mcp_agent/workflows/evaluator_optimizer/evaluator_optimizer.py +0 -481
- mcp_agent/workflows/intent_classifier/__init__.py +0 -0
- mcp_agent/workflows/intent_classifier/intent_classifier_base.py +0 -120
- mcp_agent/workflows/intent_classifier/intent_classifier_embedding.py +0 -134
- mcp_agent/workflows/intent_classifier/intent_classifier_embedding_cohere.py +0 -45
- mcp_agent/workflows/intent_classifier/intent_classifier_embedding_openai.py +0 -45
- mcp_agent/workflows/intent_classifier/intent_classifier_llm.py +0 -161
- mcp_agent/workflows/intent_classifier/intent_classifier_llm_anthropic.py +0 -60
- mcp_agent/workflows/intent_classifier/intent_classifier_llm_openai.py +0 -60
- mcp_agent/workflows/llm/__init__.py +0 -0
- mcp_agent/workflows/llm/augmented_llm.py +0 -753
- mcp_agent/workflows/llm/augmented_llm_passthrough.py +0 -241
- mcp_agent/workflows/llm/augmented_llm_playback.py +0 -109
- mcp_agent/workflows/llm/providers/__init__.py +0 -8
- mcp_agent/workflows/llm/sampling_format_converter.py +0 -22
- mcp_agent/workflows/orchestrator/__init__.py +0 -0
- mcp_agent/workflows/orchestrator/orchestrator.py +0 -578
- mcp_agent/workflows/parallel/__init__.py +0 -0
- mcp_agent/workflows/parallel/fan_in.py +0 -350
- mcp_agent/workflows/parallel/fan_out.py +0 -187
- mcp_agent/workflows/parallel/parallel_llm.py +0 -166
- mcp_agent/workflows/router/__init__.py +0 -0
- mcp_agent/workflows/router/router_base.py +0 -368
- mcp_agent/workflows/router/router_embedding.py +0 -240
- mcp_agent/workflows/router/router_embedding_cohere.py +0 -59
- mcp_agent/workflows/router/router_embedding_openai.py +0 -59
- mcp_agent/workflows/router/router_llm.py +0 -320
- mcp_agent/workflows/swarm/__init__.py +0 -0
- mcp_agent/workflows/swarm/swarm.py +0 -320
- mcp_agent/workflows/swarm/swarm_anthropic.py +0 -42
- mcp_agent/workflows/swarm/swarm_openai.py +0 -41
- {fast_agent_mcp-0.1.12.dist-info → fast_agent_mcp-0.2.0.dist-info}/WHEEL +0 -0
- {fast_agent_mcp-0.1.12.dist-info → fast_agent_mcp-0.2.0.dist-info}/entry_points.txt +0 -0
- {fast_agent_mcp-0.1.12.dist-info → fast_agent_mcp-0.2.0.dist-info}/licenses/LICENSE +0 -0
- /mcp_agent/{workflows/orchestrator → agents/workflow}/orchestrator_prompts.py +0 -0
@@ -1,78 +0,0 @@
|
|
1
|
-
"""
|
2
|
-
Parallel Workflow showing Fan Out and Fan In agents, using different models
|
3
|
-
"""
|
4
|
-
|
5
|
-
import asyncio
|
6
|
-
from mcp_agent.core.fastagent import FastAgent
|
7
|
-
|
8
|
-
# Create the application
|
9
|
-
fast = FastAgent(
|
10
|
-
"Parallel Workflow",
|
11
|
-
)
|
12
|
-
SHORT_STORY = """
|
13
|
-
The Battle of Glimmerwood
|
14
|
-
|
15
|
-
In the heart of Glimmerwood, a mystical forest knowed for its radiant trees, a small village thrived.
|
16
|
-
The villagers, who were live peacefully, shared their home with the forest's magical creatures,
|
17
|
-
especially the Glimmerfoxes whose fur shimmer like moonlight.
|
18
|
-
|
19
|
-
One fateful evening, the peace was shaterred when the infamous Dark Marauders attack.
|
20
|
-
Lead by the cunning Captain Thorn, the bandits aim to steal the precious Glimmerstones which was believed to grant immortality.
|
21
|
-
|
22
|
-
Amidst the choas, a young girl named Elara stood her ground, she rallied the villagers and devised a clever plan.
|
23
|
-
Using the forests natural defenses they lured the marauders into a trap.
|
24
|
-
As the bandits aproached the village square, a herd of Glimmerfoxes emerged, blinding them with their dazzling light,
|
25
|
-
the villagers seized the opportunity to captured the invaders.
|
26
|
-
|
27
|
-
Elara's bravery was celebrated and she was hailed as the "Guardian of Glimmerwood".
|
28
|
-
The Glimmerstones were secured in a hidden grove protected by an ancient spell.
|
29
|
-
|
30
|
-
However, not all was as it seemed. The Glimmerstones true power was never confirm,
|
31
|
-
and whispers of a hidden agenda linger among the villagers.
|
32
|
-
"""
|
33
|
-
|
34
|
-
|
35
|
-
@fast.agent(
|
36
|
-
name="proofreader",
|
37
|
-
instruction=""""Review the short story for grammar, spelling, and punctuation errors.
|
38
|
-
Identify any awkward phrasing or structural issues that could improve clarity.
|
39
|
-
Provide detailed feedback on corrections.""",
|
40
|
-
)
|
41
|
-
@fast.agent(
|
42
|
-
name="fact_checker",
|
43
|
-
instruction="""Verify the factual consistency within the story. Identify any contradictions,
|
44
|
-
logical inconsistencies, or inaccuracies in the plot, character actions, or setting.
|
45
|
-
Highlight potential issues with reasoning or coherence.""",
|
46
|
-
model="gpt-4o",
|
47
|
-
)
|
48
|
-
@fast.agent(
|
49
|
-
name="style_enforcer",
|
50
|
-
instruction="""Analyze the story for adherence to style guidelines.
|
51
|
-
Evaluate the narrative flow, clarity of expression, and tone. Suggest improvements to
|
52
|
-
enhance storytelling, readability, and engagement.""",
|
53
|
-
model="sonnet",
|
54
|
-
)
|
55
|
-
@fast.agent(
|
56
|
-
name="grader",
|
57
|
-
instruction="""Compile the feedback from the Proofreader, Fact Checker, and Style Enforcer
|
58
|
-
into a structured report. Summarize key issues and categorize them by type.
|
59
|
-
Provide actionable recommendations for improving the story,
|
60
|
-
and give an overall grade based on the feedback.""",
|
61
|
-
model="o3-mini.low",
|
62
|
-
)
|
63
|
-
@fast.parallel(
|
64
|
-
fan_out=["proofreader", "fact_checker", "style_enforcer"],
|
65
|
-
fan_in="grader",
|
66
|
-
name="parallel",
|
67
|
-
)
|
68
|
-
async def main():
|
69
|
-
# Use the app's context manager
|
70
|
-
async with fast.run() as agent:
|
71
|
-
await agent.parallel(f"student short story submission: {SHORT_STORY}")
|
72
|
-
|
73
|
-
# follow-on prompt to task agent
|
74
|
-
await agent.style_enforcer.prompt(default_prompt="STOP")
|
75
|
-
|
76
|
-
|
77
|
-
if __name__ == "__main__":
|
78
|
-
asyncio.run(main())
|
@@ -1,53 +0,0 @@
|
|
1
|
-
"""
|
2
|
-
Example MCP Agent application showing router workflow with decorator syntax.
|
3
|
-
Demonstrates router's ability to either:
|
4
|
-
1. Use tools directly to handle requests
|
5
|
-
2. Delegate requests to specialized agents
|
6
|
-
"""
|
7
|
-
|
8
|
-
import asyncio
|
9
|
-
from mcp_agent.core.fastagent import FastAgent
|
10
|
-
|
11
|
-
# Create the application
|
12
|
-
fast = FastAgent(
|
13
|
-
"Router Workflow",
|
14
|
-
)
|
15
|
-
|
16
|
-
# Sample requests demonstrating direct tool use vs agent delegation
|
17
|
-
SAMPLE_REQUESTS = [
|
18
|
-
"Download and summarize https://llmindset.co.uk/posts/2024/12/mcp-build-notes/", # Router handles directly with fetch
|
19
|
-
"Analyze the quality of the Python codebase in the current working directory", # Delegated to code expert
|
20
|
-
"What are the key principles of effective beekeeping?", # Delegated to general assistant
|
21
|
-
]
|
22
|
-
|
23
|
-
|
24
|
-
@fast.agent(
|
25
|
-
name="fetcher",
|
26
|
-
instruction="""You are an agent, with a tool enabling you to fetch URLs.""",
|
27
|
-
servers=["fetch"],
|
28
|
-
)
|
29
|
-
@fast.agent(
|
30
|
-
name="code_expert",
|
31
|
-
instruction="""You are an expert in code analysis and software engineering.
|
32
|
-
When asked about code, architecture, or development practices,
|
33
|
-
you provide thorough and practical insights.""",
|
34
|
-
servers=["filesystem"],
|
35
|
-
)
|
36
|
-
@fast.agent(
|
37
|
-
name="general_assistant",
|
38
|
-
instruction="""You are a knowledgeable assistant that provides clear,
|
39
|
-
well-reasoned responses about general topics, concepts, and principles.""",
|
40
|
-
)
|
41
|
-
@fast.router(
|
42
|
-
name="route",
|
43
|
-
model="sonnet",
|
44
|
-
agents=["code_expert", "general_assistant", "fetcher"],
|
45
|
-
)
|
46
|
-
async def main():
|
47
|
-
async with fast.run() as agent:
|
48
|
-
for request in SAMPLE_REQUESTS:
|
49
|
-
await agent.route(request)
|
50
|
-
|
51
|
-
|
52
|
-
if __name__ == "__main__":
|
53
|
-
asyncio.run(main())
|
@@ -1,23 +0,0 @@
|
|
1
|
-
# example_mcp_server.py
|
2
|
-
import asyncio
|
3
|
-
|
4
|
-
from mcp_agent.core.fastagent import FastAgent
|
5
|
-
|
6
|
-
# Define your agents as normal
|
7
|
-
fa = FastAgent("My Application")
|
8
|
-
|
9
|
-
|
10
|
-
@fa.agent("analyst", "hello, world", servers=["fetch"])
|
11
|
-
|
12
|
-
# Run the application with MCP server
|
13
|
-
async def main():
|
14
|
-
await fa.run_with_mcp_server(
|
15
|
-
transport="sse", # Use "sse" for web server, "stdio" for command line
|
16
|
-
port=8000,
|
17
|
-
server_name="MyAgents",
|
18
|
-
server_description="MCP Server exposing analyst and researcher agents",
|
19
|
-
)
|
20
|
-
|
21
|
-
|
22
|
-
if __name__ == "__main__":
|
23
|
-
asyncio.run(main())
|
mcp_agent/telemetry/__init__.py
DELETED
File without changes
|
@@ -1,18 +0,0 @@
|
|
1
|
-
import logging
|
2
|
-
from mcp_agent.config import get_settings
|
3
|
-
|
4
|
-
logger = logging.getLogger(__name__)
|
5
|
-
|
6
|
-
|
7
|
-
def send_usage_data():
|
8
|
-
config = get_settings()
|
9
|
-
if not config.usage_telemetry.enabled:
|
10
|
-
logger.info("Usage tracking is disabled")
|
11
|
-
return
|
12
|
-
|
13
|
-
# TODO: saqadri - implement usage tracking
|
14
|
-
# data = {"installation_id": str(uuid.uuid4()), "version": "0.1.0"}
|
15
|
-
# try:
|
16
|
-
# requests.post("https://telemetry.example.com/usage", json=data, timeout=2)
|
17
|
-
# except:
|
18
|
-
# pass
|
mcp_agent/workflows/__init__.py
DELETED
File without changes
|
File without changes
|
@@ -1,61 +0,0 @@
|
|
1
|
-
from abc import ABC, abstractmethod
|
2
|
-
from typing import Dict, List
|
3
|
-
|
4
|
-
from numpy import float32
|
5
|
-
from numpy.typing import NDArray
|
6
|
-
from sklearn.metrics.pairwise import cosine_similarity
|
7
|
-
|
8
|
-
from mcp_agent.context_dependent import ContextDependent
|
9
|
-
|
10
|
-
|
11
|
-
FloatArray = NDArray[float32]
|
12
|
-
|
13
|
-
|
14
|
-
class EmbeddingModel(ABC, ContextDependent):
|
15
|
-
"""Abstract interface for embedding models"""
|
16
|
-
|
17
|
-
@abstractmethod
|
18
|
-
async def embed(self, data: List[str]) -> FloatArray:
|
19
|
-
"""
|
20
|
-
Generate embeddings for a list of messages
|
21
|
-
|
22
|
-
Args:
|
23
|
-
data: List of text strings to embed
|
24
|
-
|
25
|
-
Returns:
|
26
|
-
Array of embeddings, shape (len(texts), embedding_dim)
|
27
|
-
"""
|
28
|
-
|
29
|
-
@property
|
30
|
-
@abstractmethod
|
31
|
-
def embedding_dim(self) -> int:
|
32
|
-
"""Return the dimensionality of the embeddings"""
|
33
|
-
|
34
|
-
|
35
|
-
def compute_similarity_scores(
|
36
|
-
embedding_a: FloatArray, embedding_b: FloatArray
|
37
|
-
) -> Dict[str, float]:
|
38
|
-
"""
|
39
|
-
Compute different similarity metrics between embeddings
|
40
|
-
"""
|
41
|
-
# Reshape for sklearn's cosine_similarity
|
42
|
-
a_emb = embedding_a.reshape(1, -1)
|
43
|
-
b_emb = embedding_b.reshape(1, -1)
|
44
|
-
|
45
|
-
cosine_sim = float(cosine_similarity(a_emb, b_emb)[0, 0])
|
46
|
-
|
47
|
-
# Could add other similarity metrics here
|
48
|
-
return {
|
49
|
-
"cosine": cosine_sim,
|
50
|
-
# "euclidean": float(euclidean_similarity),
|
51
|
-
# "dot_product": float(dot_product)
|
52
|
-
}
|
53
|
-
|
54
|
-
|
55
|
-
def compute_confidence(similarity_scores: Dict[str, float]) -> float:
|
56
|
-
"""
|
57
|
-
Compute overall confidence score from individual similarity metrics
|
58
|
-
"""
|
59
|
-
# For now, just use cosine similarity as confidence
|
60
|
-
# Could implement more sophisticated combination of scores
|
61
|
-
return similarity_scores["cosine"]
|
@@ -1,49 +0,0 @@
|
|
1
|
-
from typing import List, Optional, TYPE_CHECKING
|
2
|
-
|
3
|
-
from cohere import Client
|
4
|
-
from numpy import array, float32
|
5
|
-
|
6
|
-
from mcp_agent.workflows.embedding.embedding_base import EmbeddingModel, FloatArray
|
7
|
-
|
8
|
-
if TYPE_CHECKING:
|
9
|
-
from mcp_agent.context import Context
|
10
|
-
|
11
|
-
|
12
|
-
class CohereEmbeddingModel(EmbeddingModel):
|
13
|
-
"""Cohere embedding model implementation"""
|
14
|
-
|
15
|
-
def __init__(
|
16
|
-
self,
|
17
|
-
model: str = "embed-multilingual-v3.0",
|
18
|
-
context: Optional["Context"] = None,
|
19
|
-
**kwargs,
|
20
|
-
):
|
21
|
-
super().__init__(context=context, **kwargs)
|
22
|
-
self.client = Client(api_key=self.context.config.cohere.api_key)
|
23
|
-
self.model = model
|
24
|
-
# Cache the dimension since it's fixed per model
|
25
|
-
# https://docs.cohere.com/v2/docs/cohere-embed
|
26
|
-
self._embedding_dim = {
|
27
|
-
"embed-english-v2.0": 4096,
|
28
|
-
"embed-english-light-v2.0": 1024,
|
29
|
-
"embed-english-v3.0": 1024,
|
30
|
-
"embed-english-light-v3.0": 384,
|
31
|
-
"embed-multilingual-v2.0": 768,
|
32
|
-
"embed-multilingual-v3.0": 1024,
|
33
|
-
"embed-multilingual-light-v3.0": 384,
|
34
|
-
}[model]
|
35
|
-
|
36
|
-
async def embed(self, data: List[str]) -> FloatArray:
|
37
|
-
response = self.client.embed(
|
38
|
-
texts=data,
|
39
|
-
model=self.model,
|
40
|
-
input_type="classification",
|
41
|
-
embedding_types=["float"],
|
42
|
-
)
|
43
|
-
|
44
|
-
embeddings = array(response.embeddings, dtype=float32)
|
45
|
-
return embeddings
|
46
|
-
|
47
|
-
@property
|
48
|
-
def embedding_dim(self) -> int:
|
49
|
-
return self._embedding_dim
|
@@ -1,46 +0,0 @@
|
|
1
|
-
from typing import List, Optional, TYPE_CHECKING
|
2
|
-
|
3
|
-
from numpy import array, float32, stack
|
4
|
-
from openai import OpenAI
|
5
|
-
|
6
|
-
from mcp_agent.workflows.embedding.embedding_base import EmbeddingModel, FloatArray
|
7
|
-
|
8
|
-
if TYPE_CHECKING:
|
9
|
-
from mcp_agent.context import Context
|
10
|
-
|
11
|
-
|
12
|
-
class OpenAIEmbeddingModel(EmbeddingModel):
|
13
|
-
"""OpenAI embedding model implementation"""
|
14
|
-
|
15
|
-
def __init__(
|
16
|
-
self, model: str = "text-embedding-3-small", context: Optional["Context"] = None
|
17
|
-
):
|
18
|
-
super().__init__(context=context)
|
19
|
-
self.client = OpenAI(api_key=self.context.config.openai.api_key)
|
20
|
-
self.model = model
|
21
|
-
# Cache the dimension since it's fixed per model
|
22
|
-
self._embedding_dim = {
|
23
|
-
"text-embedding-3-small": 1536,
|
24
|
-
"text-embedding-3-large": 3072,
|
25
|
-
}[model]
|
26
|
-
|
27
|
-
async def embed(self, data: List[str]) -> FloatArray:
|
28
|
-
response = self.client.embeddings.create(
|
29
|
-
model=self.model, input=data, encoding_format="float"
|
30
|
-
)
|
31
|
-
|
32
|
-
# Sort the embeddings by their index to ensure correct order
|
33
|
-
sorted_embeddings = sorted(response.data, key=lambda x: x["index"])
|
34
|
-
|
35
|
-
# Stack all embeddings into a single array
|
36
|
-
embeddings = stack(
|
37
|
-
[
|
38
|
-
array(embedding["embedding"], dtype=float32)
|
39
|
-
for embedding in sorted_embeddings
|
40
|
-
]
|
41
|
-
)
|
42
|
-
return embeddings
|
43
|
-
|
44
|
-
@property
|
45
|
-
def embedding_dim(self) -> int:
|
46
|
-
return self._embedding_dim
|
File without changes
|