agentic-ai-kit 0.1.1__tar.gz

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.
Files changed (50) hide show
  1. agentic_ai_kit-0.1.1/PKG-INFO +201 -0
  2. agentic_ai_kit-0.1.1/README.md +172 -0
  3. agentic_ai_kit-0.1.1/agentic_ai/__init__.py +24 -0
  4. agentic_ai_kit-0.1.1/agentic_ai/agents/__init__.py +17 -0
  5. agentic_ai_kit-0.1.1/agentic_ai/agents/analyst_agent.py +239 -0
  6. agentic_ai_kit-0.1.1/agentic_ai/agents/auto_model_agent.py +156 -0
  7. agentic_ai_kit-0.1.1/agentic_ai/agents/base.py +280 -0
  8. agentic_ai_kit-0.1.1/agentic_ai/agents/json_agent.py +73 -0
  9. agentic_ai_kit-0.1.1/agentic_ai/agents/mle_agent.py +590 -0
  10. agentic_ai_kit-0.1.1/agentic_ai/agents/reasoning_agent.py +97 -0
  11. agentic_ai_kit-0.1.1/agentic_ai/agents/tool_agent.py +171 -0
  12. agentic_ai_kit-0.1.1/agentic_ai/examples/01_talking_agents.py +9 -0
  13. agentic_ai_kit-0.1.1/agentic_ai/examples/02_personality_and_memory.py +12 -0
  14. agentic_ai_kit-0.1.1/agentic_ai/examples/03_tool_agent_weather.py +10 -0
  15. agentic_ai_kit-0.1.1/agentic_ai/examples/04_json_agent.py +11 -0
  16. agentic_ai_kit-0.1.1/agentic_ai/examples/05_reasoning_agent.py +9 -0
  17. agentic_ai_kit-0.1.1/agentic_ai/examples/06_supervisor.py +16 -0
  18. agentic_ai_kit-0.1.1/agentic_ai/examples/07_parallel_agents.py +15 -0
  19. agentic_ai_kit-0.1.1/agentic_ai/examples/08_debate.py +16 -0
  20. agentic_ai_kit-0.1.1/agentic_ai/examples/09_long_term_memory.py +23 -0
  21. agentic_ai_kit-0.1.1/agentic_ai/examples/10_rag_agent.py +29 -0
  22. agentic_ai_kit-0.1.1/agentic_ai/examples/11_mcp_demo.py +41 -0
  23. agentic_ai_kit-0.1.1/agentic_ai/examples/__init__.py +1 -0
  24. agentic_ai_kit-0.1.1/agentic_ai/mcp/__init__.py +4 -0
  25. agentic_ai_kit-0.1.1/agentic_ai/mcp/client.py +55 -0
  26. agentic_ai_kit-0.1.1/agentic_ai/mcp/server.py +133 -0
  27. agentic_ai_kit-0.1.1/agentic_ai/memory/__init__.py +5 -0
  28. agentic_ai_kit-0.1.1/agentic_ai/memory/long_term.py +104 -0
  29. agentic_ai_kit-0.1.1/agentic_ai/memory/shared.py +52 -0
  30. agentic_ai_kit-0.1.1/agentic_ai/memory/short_term.py +43 -0
  31. agentic_ai_kit-0.1.1/agentic_ai/patterns/__init__.py +5 -0
  32. agentic_ai_kit-0.1.1/agentic_ai/patterns/debate.py +49 -0
  33. agentic_ai_kit-0.1.1/agentic_ai/patterns/orchestrator.py +91 -0
  34. agentic_ai_kit-0.1.1/agentic_ai/patterns/parallel.py +39 -0
  35. agentic_ai_kit-0.1.1/agentic_ai/py.typed +0 -0
  36. agentic_ai_kit-0.1.1/agentic_ai/rag/__init__.py +6 -0
  37. agentic_ai_kit-0.1.1/agentic_ai/rag/chunker.py +23 -0
  38. agentic_ai_kit-0.1.1/agentic_ai/rag/embedder.py +30 -0
  39. agentic_ai_kit-0.1.1/agentic_ai/rag/rag_agent.py +42 -0
  40. agentic_ai_kit-0.1.1/agentic_ai/rag/vector_store.py +51 -0
  41. agentic_ai_kit-0.1.1/agentic_ai/tools/__init__.py +8 -0
  42. agentic_ai_kit-0.1.1/agentic_ai/tools/builtins.py +51 -0
  43. agentic_ai_kit-0.1.1/agentic_ai/tools/registry.py +120 -0
  44. agentic_ai_kit-0.1.1/agentic_ai_kit.egg-info/PKG-INFO +201 -0
  45. agentic_ai_kit-0.1.1/agentic_ai_kit.egg-info/SOURCES.txt +48 -0
  46. agentic_ai_kit-0.1.1/agentic_ai_kit.egg-info/dependency_links.txt +1 -0
  47. agentic_ai_kit-0.1.1/agentic_ai_kit.egg-info/requires.txt +10 -0
  48. agentic_ai_kit-0.1.1/agentic_ai_kit.egg-info/top_level.txt +1 -0
  49. agentic_ai_kit-0.1.1/pyproject.toml +44 -0
  50. agentic_ai_kit-0.1.1/setup.cfg +4 -0
@@ -0,0 +1,201 @@
1
+ Metadata-Version: 2.4
2
+ Name: agentic-ai-kit
3
+ Version: 0.1.1
4
+ Summary: A progressive framework for building agentic AI systems with Google Gemini.
5
+ Author-email: Shashank Singh <shashank14581@gmail.com>
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/shashank14581/agentic-ai-kit
8
+ Project-URL: Repository, https://github.com/shashank14581/agentic-ai-kit
9
+ Keywords: agentic-ai,llm,gemini,agents,rag,mcp
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Programming Language :: Python :: 3.10
14
+ Classifier: Programming Language :: Python :: 3.11
15
+ Classifier: Programming Language :: Python :: 3.12
16
+ Classifier: License :: OSI Approved :: MIT License
17
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
18
+ Requires-Python: >=3.10
19
+ Description-Content-Type: text/markdown
20
+ Requires-Dist: google-genai>=1.0.0
21
+ Requires-Dist: requests>=2.31.0
22
+ Requires-Dist: numpy>=1.24.0
23
+ Requires-Dist: pandas>=2.0.0
24
+ Requires-Dist: scikit-learn>=1.3.0
25
+ Provides-Extra: dev
26
+ Requires-Dist: pytest; extra == "dev"
27
+ Requires-Dist: build; extra == "dev"
28
+ Requires-Dist: twine; extra == "dev"
29
+
30
+ # agentic-ai
31
+
32
+ A progressive Python framework for building **agentic AI systems with Google Gemini**.
33
+
34
+ `agentic-ai` turns common agentic AI patterns into reusable Python components: talking agents, memory, tool use, structured JSON output, reasoning workflows, supervisor delegation, parallel execution, debate, RAG, MCP tool servers, and applied data science agents.
35
+
36
+ This project started with a simple question: what happens if two agents talk and one of them is a dog that can bark? From that playful starting point, the patterns evolved into a reusable framework for building composable agentic systems.
37
+
38
+ ---
39
+
40
+ ## Installation
41
+
42
+ ```bash
43
+ pip install agentic-ai
44
+ ```
45
+
46
+ For the latest GitHub version:
47
+
48
+ ```bash
49
+ pip install --upgrade --force-reinstall git+https://github.com/shashank14581/agentic-ai-kit.git
50
+ ```
51
+
52
+ Set your Gemini API key:
53
+
54
+ ```bash
55
+ export GEMINI_API_KEY="your-key-here"
56
+ ```
57
+
58
+ In Google Colab, store `GEMINI_API_KEY` in Colab Secrets and load it into `os.environ`.
59
+
60
+ ---
61
+
62
+ ## Quick start
63
+
64
+ ```python
65
+ from agentic_ai.agents import BaseAgent
66
+
67
+ agent = BaseAgent(
68
+ name="Alfred",
69
+ sys_prompt="You are a witty British butler.",
70
+ model="gemini-1.5-flash",
71
+ extract_memory=False,
72
+ )
73
+
74
+ agent.think("Good morning! What should I do today?", stream=False)
75
+ ```
76
+
77
+ ---
78
+
79
+ ## What it supports
80
+
81
+ | Capability | Component |
82
+ |---|---|
83
+ | Simple conversational agents | `BaseAgent` |
84
+ | Persona and short-term memory | `BaseAgent`, `ShortTermMemory` |
85
+ | LLM-extracted durable facts | `BaseAgent.facts_store` |
86
+ | SQLite-backed long-term memory | `LongTermMemory` |
87
+ | Shared state across agents | `SharedMemory` |
88
+ | Gemini function-calling tools | `ToolAgent`, `ToolRegistry` |
89
+ | Structured JSON output | `JsonAgent` |
90
+ | Reasoning / planning workflows | `ReasoningAgent` |
91
+ | Round-robin conversations | `run_conversation` |
92
+ | Supervisor delegation | `run_supervisor` |
93
+ | Parallel fan-out execution | `run_parallel` |
94
+ | Debate and judge pattern | `run_debate` |
95
+ | Chunking, embeddings, vector search | `chunk_text`, `embed_texts`, `VectorStore` |
96
+ | Retrieval-augmented generation | `RAGAgent` |
97
+ | MCP-style tool server and client | `MCPServer`, `MCPClient` |
98
+ | DataFrame analysis | `AnalystAgent` |
99
+ | Machine learning model training | `MLEAgent` |
100
+ | End-to-end auto modeling workflow | `AutoModelAgent` |
101
+
102
+ ---
103
+
104
+ ## Package layout
105
+
106
+ ```text
107
+ agentic_ai/
108
+ ├── agents/
109
+ │ ├── base.py # BaseAgent — memory, streaming, durable facts
110
+ │ ├── tool_agent.py # ToolAgent — Gemini function-calling
111
+ │ ├── json_agent.py # JsonAgent — structured JSON responses
112
+ │ ├── reasoning_agent.py # ReasoningAgent — planning-style workflows
113
+ │ ├── analyst_agent.py # AnalystAgent — dataframe profiling and analysis
114
+ │ ├── mle_agent.py # MLEAgent — train and compare sklearn models
115
+ │ └── auto_model_agent.py # AutoModelAgent — profile data, train model, summarize results
116
+ ├── memory/
117
+ │ ├── short_term.py # Sliding-window turn memory
118
+ │ ├── long_term.py # SQLite-backed persistent memory
119
+ │ └── shared.py # Thread-safe shared state across agents
120
+ ├── tools/
121
+ │ ├── registry.py # @tool decorator and ToolRegistry
122
+ │ └── builtins.py # get_weather, add_numbers, search_wikipedia
123
+ ├── patterns/
124
+ │ ├── orchestrator.py # run_conversation, run_supervisor
125
+ │ ├── parallel.py # run_parallel fan-out execution
126
+ │ └── debate.py # run_debate with debaters and judge
127
+ ├── rag/
128
+ │ ├── chunker.py # chunk_text
129
+ │ ├── embedder.py # embed_texts using Gemini embeddings
130
+ │ ├── vector_store.py # In-memory cosine-similarity vector store
131
+ │ └── rag_agent.py # RAGAgent ingest → retrieve → generate
132
+ └── mcp/
133
+ ├── server.py # MCPServer — HTTP tool server
134
+ └── client.py # MCPClient — calls a running MCPServer
135
+ ```
136
+
137
+ ---
138
+
139
+ ## Examples
140
+
141
+ All examples live in `agentic_ai/examples/` and can be run directly:
142
+
143
+ | File | Concept |
144
+ |---|---|
145
+ | `01_talking_agents.py` | Two agents in a round-robin conversation |
146
+ | `02_personality_and_memory.py` | Agent with persona and short-term memory |
147
+ | `03_tool_agent_weather.py` | Real weather tool through Gemini function-calling |
148
+ | `04_json_agent.py` | Structured JSON output |
149
+ | `05_reasoning_agent.py` | Planning-style reasoning workflow |
150
+ | `06_supervisor.py` | Supervisor delegates to specialist workers |
151
+ | `07_parallel_agents.py` | Fan-out: same task, multiple agents in parallel |
152
+ | `08_debate.py` | Two agents debate; a judge picks a winner |
153
+ | `09_long_term_memory.py` | SQLite-backed memory across sessions |
154
+ | `10_rag_agent.py` | Document ingestion and retrieval-augmented generation |
155
+ | `11_mcp_demo.py` | MCP tool server and agent integration |
156
+
157
+ ---
158
+
159
+ ## Data science agents
160
+
161
+ The kit also includes applied analytics and modeling agents.
162
+
163
+ ```python
164
+ from agentic_ai.agents import AutoModelAgent
165
+
166
+ agent = AutoModelAgent(
167
+ project_context="Marketing audience selection and customer conversion modeling.",
168
+ model="gemini-1.5-flash",
169
+ )
170
+
171
+ result = agent.run(
172
+ df=df,
173
+ target_col="converted",
174
+ objective="Predict customer conversion.",
175
+ interpret=False,
176
+ )
177
+
178
+ print(result["best_model"])
179
+ print(result["best_score"])
180
+ ```
181
+
182
+ Use `interpret=False` when you want local sklearn training without an additional Gemini interpretation call.
183
+
184
+ ---
185
+
186
+ ## Design philosophy
187
+
188
+ The framework separates **reasoning** from **execution**.
189
+
190
+ - Agents can reason, plan, route, critique, and summarize.
191
+ - Utilities can execute deterministic work such as memory lookup, chunking, vector search, dataframe profiling, and model training.
192
+ - Pattern functions provide lightweight orchestration.
193
+ - Specialist agents provide higher-level user-facing workflows.
194
+
195
+ This makes the library useful for both learning agentic AI patterns and building practical data science or automation workflows.
196
+
197
+ ---
198
+
199
+ ## License
200
+
201
+ MIT
@@ -0,0 +1,172 @@
1
+ # agentic-ai
2
+
3
+ A progressive Python framework for building **agentic AI systems with Google Gemini**.
4
+
5
+ `agentic-ai` turns common agentic AI patterns into reusable Python components: talking agents, memory, tool use, structured JSON output, reasoning workflows, supervisor delegation, parallel execution, debate, RAG, MCP tool servers, and applied data science agents.
6
+
7
+ This project started with a simple question: what happens if two agents talk and one of them is a dog that can bark? From that playful starting point, the patterns evolved into a reusable framework for building composable agentic systems.
8
+
9
+ ---
10
+
11
+ ## Installation
12
+
13
+ ```bash
14
+ pip install agentic-ai
15
+ ```
16
+
17
+ For the latest GitHub version:
18
+
19
+ ```bash
20
+ pip install --upgrade --force-reinstall git+https://github.com/shashank14581/agentic-ai-kit.git
21
+ ```
22
+
23
+ Set your Gemini API key:
24
+
25
+ ```bash
26
+ export GEMINI_API_KEY="your-key-here"
27
+ ```
28
+
29
+ In Google Colab, store `GEMINI_API_KEY` in Colab Secrets and load it into `os.environ`.
30
+
31
+ ---
32
+
33
+ ## Quick start
34
+
35
+ ```python
36
+ from agentic_ai.agents import BaseAgent
37
+
38
+ agent = BaseAgent(
39
+ name="Alfred",
40
+ sys_prompt="You are a witty British butler.",
41
+ model="gemini-1.5-flash",
42
+ extract_memory=False,
43
+ )
44
+
45
+ agent.think("Good morning! What should I do today?", stream=False)
46
+ ```
47
+
48
+ ---
49
+
50
+ ## What it supports
51
+
52
+ | Capability | Component |
53
+ |---|---|
54
+ | Simple conversational agents | `BaseAgent` |
55
+ | Persona and short-term memory | `BaseAgent`, `ShortTermMemory` |
56
+ | LLM-extracted durable facts | `BaseAgent.facts_store` |
57
+ | SQLite-backed long-term memory | `LongTermMemory` |
58
+ | Shared state across agents | `SharedMemory` |
59
+ | Gemini function-calling tools | `ToolAgent`, `ToolRegistry` |
60
+ | Structured JSON output | `JsonAgent` |
61
+ | Reasoning / planning workflows | `ReasoningAgent` |
62
+ | Round-robin conversations | `run_conversation` |
63
+ | Supervisor delegation | `run_supervisor` |
64
+ | Parallel fan-out execution | `run_parallel` |
65
+ | Debate and judge pattern | `run_debate` |
66
+ | Chunking, embeddings, vector search | `chunk_text`, `embed_texts`, `VectorStore` |
67
+ | Retrieval-augmented generation | `RAGAgent` |
68
+ | MCP-style tool server and client | `MCPServer`, `MCPClient` |
69
+ | DataFrame analysis | `AnalystAgent` |
70
+ | Machine learning model training | `MLEAgent` |
71
+ | End-to-end auto modeling workflow | `AutoModelAgent` |
72
+
73
+ ---
74
+
75
+ ## Package layout
76
+
77
+ ```text
78
+ agentic_ai/
79
+ ├── agents/
80
+ │ ├── base.py # BaseAgent — memory, streaming, durable facts
81
+ │ ├── tool_agent.py # ToolAgent — Gemini function-calling
82
+ │ ├── json_agent.py # JsonAgent — structured JSON responses
83
+ │ ├── reasoning_agent.py # ReasoningAgent — planning-style workflows
84
+ │ ├── analyst_agent.py # AnalystAgent — dataframe profiling and analysis
85
+ │ ├── mle_agent.py # MLEAgent — train and compare sklearn models
86
+ │ └── auto_model_agent.py # AutoModelAgent — profile data, train model, summarize results
87
+ ├── memory/
88
+ │ ├── short_term.py # Sliding-window turn memory
89
+ │ ├── long_term.py # SQLite-backed persistent memory
90
+ │ └── shared.py # Thread-safe shared state across agents
91
+ ├── tools/
92
+ │ ├── registry.py # @tool decorator and ToolRegistry
93
+ │ └── builtins.py # get_weather, add_numbers, search_wikipedia
94
+ ├── patterns/
95
+ │ ├── orchestrator.py # run_conversation, run_supervisor
96
+ │ ├── parallel.py # run_parallel fan-out execution
97
+ │ └── debate.py # run_debate with debaters and judge
98
+ ├── rag/
99
+ │ ├── chunker.py # chunk_text
100
+ │ ├── embedder.py # embed_texts using Gemini embeddings
101
+ │ ├── vector_store.py # In-memory cosine-similarity vector store
102
+ │ └── rag_agent.py # RAGAgent ingest → retrieve → generate
103
+ └── mcp/
104
+ ├── server.py # MCPServer — HTTP tool server
105
+ └── client.py # MCPClient — calls a running MCPServer
106
+ ```
107
+
108
+ ---
109
+
110
+ ## Examples
111
+
112
+ All examples live in `agentic_ai/examples/` and can be run directly:
113
+
114
+ | File | Concept |
115
+ |---|---|
116
+ | `01_talking_agents.py` | Two agents in a round-robin conversation |
117
+ | `02_personality_and_memory.py` | Agent with persona and short-term memory |
118
+ | `03_tool_agent_weather.py` | Real weather tool through Gemini function-calling |
119
+ | `04_json_agent.py` | Structured JSON output |
120
+ | `05_reasoning_agent.py` | Planning-style reasoning workflow |
121
+ | `06_supervisor.py` | Supervisor delegates to specialist workers |
122
+ | `07_parallel_agents.py` | Fan-out: same task, multiple agents in parallel |
123
+ | `08_debate.py` | Two agents debate; a judge picks a winner |
124
+ | `09_long_term_memory.py` | SQLite-backed memory across sessions |
125
+ | `10_rag_agent.py` | Document ingestion and retrieval-augmented generation |
126
+ | `11_mcp_demo.py` | MCP tool server and agent integration |
127
+
128
+ ---
129
+
130
+ ## Data science agents
131
+
132
+ The kit also includes applied analytics and modeling agents.
133
+
134
+ ```python
135
+ from agentic_ai.agents import AutoModelAgent
136
+
137
+ agent = AutoModelAgent(
138
+ project_context="Marketing audience selection and customer conversion modeling.",
139
+ model="gemini-1.5-flash",
140
+ )
141
+
142
+ result = agent.run(
143
+ df=df,
144
+ target_col="converted",
145
+ objective="Predict customer conversion.",
146
+ interpret=False,
147
+ )
148
+
149
+ print(result["best_model"])
150
+ print(result["best_score"])
151
+ ```
152
+
153
+ Use `interpret=False` when you want local sklearn training without an additional Gemini interpretation call.
154
+
155
+ ---
156
+
157
+ ## Design philosophy
158
+
159
+ The framework separates **reasoning** from **execution**.
160
+
161
+ - Agents can reason, plan, route, critique, and summarize.
162
+ - Utilities can execute deterministic work such as memory lookup, chunking, vector search, dataframe profiling, and model training.
163
+ - Pattern functions provide lightweight orchestration.
164
+ - Specialist agents provide higher-level user-facing workflows.
165
+
166
+ This makes the library useful for both learning agentic AI patterns and building practical data science or automation workflows.
167
+
168
+ ---
169
+
170
+ ## License
171
+
172
+ MIT
@@ -0,0 +1,24 @@
1
+ """
2
+ agentic_ai — A progressive framework for building agentic AI systems with Google Gemini.
3
+ """
4
+
5
+ from agentic_ai.agents.base import BaseAgent
6
+ from agentic_ai.agents.tool_agent import ToolAgent
7
+ from agentic_ai.agents.json_agent import JsonAgent
8
+ from agentic_ai.agents.reasoning_agent import ReasoningAgent
9
+ from agentic_ai.memory.short_term import ShortTermMemory
10
+ from agentic_ai.memory.long_term import LongTermMemory
11
+ from agentic_ai.memory.shared import SharedMemory
12
+ from agentic_ai.agents.auto_model_agent import AutoModelAgent
13
+
14
+ __version__ = "0.1.0"
15
+ __all__ = [
16
+ "BaseAgent",
17
+ "ToolAgent",
18
+ "JsonAgent",
19
+ "ReasoningAgent",
20
+ "ShortTermMemory",
21
+ "LongTermMemory",
22
+ "SharedMemory",
23
+ "AutoModelAgent",
24
+ ]
@@ -0,0 +1,17 @@
1
+ from agentic_ai.agents.base import BaseAgent
2
+ from agentic_ai.agents.tool_agent import ToolAgent
3
+ from agentic_ai.agents.json_agent import JsonAgent
4
+ from agentic_ai.agents.reasoning_agent import ReasoningAgent
5
+ from agentic_ai.agents.analyst_agent import AnalystAgent
6
+ from agentic_ai.agents.mle_agent import MLEAgent
7
+ from agentic_ai.agents.auto_model_agent import AutoModelAgent
8
+
9
+ __all__ = [
10
+ "BaseAgent",
11
+ "ToolAgent",
12
+ "JsonAgent",
13
+ "ReasoningAgent",
14
+ "AnalystAgent",
15
+ "MLEAgent",
16
+ "AutoModelAgent",
17
+ ]
@@ -0,0 +1,239 @@
1
+ """
2
+ AnalystAgent: dataframe-aware analytics specialist agent.
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ from typing import Any
8
+
9
+ import pandas as pd
10
+
11
+ from agentic_ai.agents.base import BaseAgent
12
+
13
+
14
+ _ANALYST_PROMPT = """
15
+ You are an elite Decision Scientist.
16
+
17
+ You help users analyze pandas DataFrames and business datasets.
18
+
19
+ You are expected to reason using:
20
+ - dataframe schema
21
+ - row counts
22
+ - missing values
23
+ - numeric summaries
24
+ - categorical distributions
25
+ - group-level metrics
26
+ - data quality risks
27
+ - business interpretation
28
+
29
+ Always separate:
30
+ 1. what the data says
31
+ 2. what it might mean
32
+ 3. what should be checked next
33
+
34
+ Preferred response structure:
35
+
36
+ BUSINESS QUESTION
37
+ DATA READ
38
+ KEY OBSERVATIONS
39
+ DATA QUALITY RISKS
40
+ RECOMMENDED ANALYSIS
41
+ DECISION SUMMARY
42
+
43
+ Do not invent numbers.
44
+ If dataframe context is provided, use it explicitly.
45
+ """.strip()
46
+
47
+
48
+ class AnalystAgent(BaseAgent):
49
+ """Analytics-focused specialist agent with dataframe tools."""
50
+
51
+ def __init__(
52
+ self,
53
+ name: str = "Analyst",
54
+ domain_context: str | None = None,
55
+ model: str = "gemini-2.5-flash-lite",
56
+ api_key: str | None = None,
57
+ memory_window: int = 3,
58
+ max_turns: int | None = None,
59
+ thinking_budget: int = 0,
60
+ ):
61
+ prompt = _ANALYST_PROMPT
62
+
63
+ if domain_context:
64
+ prompt += f"\n\nDOMAIN CONTEXT:\n{domain_context.strip()}"
65
+
66
+ super().__init__(
67
+ name=name,
68
+ sys_prompt=prompt,
69
+ model=model,
70
+ api_key=api_key,
71
+ memory_window=memory_window,
72
+ max_turns=max_turns,
73
+ thinking_budget=thinking_budget,
74
+ )
75
+
76
+ # ------------------------------------------------------------------
77
+ # DataFrame tools
78
+ # ------------------------------------------------------------------
79
+
80
+ def profile_dataframe(self, df: pd.DataFrame) -> dict[str, Any]:
81
+ """Return a consistent dataframe profile."""
82
+
83
+ self._validate_dataframe(df)
84
+
85
+ return {
86
+ "rows": int(df.shape[0]),
87
+ "columns": int(df.shape[1]),
88
+ "column_names": list(df.columns),
89
+ "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},
90
+ "missing_values": df.isna().sum().to_dict(),
91
+ "missing_percentage": (
92
+ df.isna().mean().mul(100).round(2).to_dict()
93
+ ),
94
+ "duplicate_rows": int(df.duplicated().sum()),
95
+ "numeric_columns": list(df.select_dtypes(include="number").columns),
96
+ "categorical_columns": list(
97
+ df.select_dtypes(include=["object", "category", "bool"]).columns
98
+ ),
99
+ "datetime_columns": list(
100
+ df.select_dtypes(include=["datetime", "datetimetz"]).columns
101
+ ),
102
+ }
103
+
104
+ def numeric_summary(self, df: pd.DataFrame) -> dict[str, Any]:
105
+ """Return numeric column summary."""
106
+
107
+ self._validate_dataframe(df)
108
+
109
+ numeric_df = df.select_dtypes(include="number")
110
+
111
+ if numeric_df.empty:
112
+ return {"message": "No numeric columns found."}
113
+
114
+ return numeric_df.describe().round(4).to_dict()
115
+
116
+ def categorical_summary(
117
+ self,
118
+ df: pd.DataFrame,
119
+ top_n: int = 10,
120
+ ) -> dict[str, Any]:
121
+ """Return top category counts for categorical columns."""
122
+
123
+ self._validate_dataframe(df)
124
+
125
+ categorical_df = df.select_dtypes(include=["object", "category", "bool"])
126
+
127
+ if categorical_df.empty:
128
+ return {"message": "No categorical columns found."}
129
+
130
+ summary = {}
131
+
132
+ for col in categorical_df.columns:
133
+ summary[col] = (
134
+ categorical_df[col]
135
+ .value_counts(dropna=False)
136
+ .head(top_n)
137
+ .to_dict()
138
+ )
139
+
140
+ return summary
141
+
142
+ def groupby_summary(
143
+ self,
144
+ df: pd.DataFrame,
145
+ group_col: str,
146
+ metric_col: str,
147
+ agg: str = "mean",
148
+ ) -> dict[str, Any]:
149
+ """Return grouped metric summary."""
150
+
151
+ self._validate_dataframe(df)
152
+ self._validate_columns(df, [group_col, metric_col])
153
+
154
+ allowed_aggs = {"mean", "sum", "count", "median", "min", "max"}
155
+
156
+ if agg not in allowed_aggs:
157
+ raise ValueError(f"agg must be one of {sorted(allowed_aggs)}")
158
+
159
+ result = (
160
+ df.groupby(group_col, dropna=False)[metric_col]
161
+ .agg(agg)
162
+ .reset_index()
163
+ .sort_values(metric_col, ascending=False)
164
+ )
165
+
166
+ return result.to_dict(orient="records")
167
+
168
+ def correlation_summary(self, df: pd.DataFrame) -> dict[str, Any]:
169
+ """Return numeric correlation matrix."""
170
+
171
+ self._validate_dataframe(df)
172
+
173
+ numeric_df = df.select_dtypes(include="number")
174
+
175
+ if numeric_df.shape[1] < 2:
176
+ return {"message": "At least two numeric columns are needed."}
177
+
178
+ return numeric_df.corr().round(4).to_dict()
179
+
180
+ def dataframe_context(
181
+ self,
182
+ df: pd.DataFrame,
183
+ include_numeric: bool = True,
184
+ include_categorical: bool = True,
185
+ ) -> str:
186
+ """Create dataframe context for the LLM."""
187
+
188
+ context = {
189
+ "profile": self.profile_dataframe(df),
190
+ }
191
+
192
+ if include_numeric:
193
+ context["numeric_summary"] = self.numeric_summary(df)
194
+
195
+ if include_categorical:
196
+ context["categorical_summary"] = self.categorical_summary(df)
197
+
198
+ return str(context)
199
+
200
+ def analyze_dataframe(
201
+ self,
202
+ df: pd.DataFrame,
203
+ question: str,
204
+ stream: bool = True,
205
+ ) -> str:
206
+ """Analyze a dataframe using generated dataframe context."""
207
+
208
+ context = self.dataframe_context(df)
209
+
210
+ prompt = f"""
211
+ Analyze the dataframe for the following question.
212
+
213
+ QUESTION:
214
+ {question}
215
+
216
+ DATAFRAME CONTEXT:
217
+ {context}
218
+ """.strip()
219
+
220
+ return self.think(prompt, stream=stream)
221
+
222
+ # ------------------------------------------------------------------
223
+ # Validation helpers
224
+ # ------------------------------------------------------------------
225
+
226
+ @staticmethod
227
+ def _validate_dataframe(df: pd.DataFrame) -> None:
228
+ if not isinstance(df, pd.DataFrame):
229
+ raise TypeError("Expected a pandas DataFrame.")
230
+
231
+ if df.empty:
232
+ raise ValueError("DataFrame is empty.")
233
+
234
+ @staticmethod
235
+ def _validate_columns(df: pd.DataFrame, columns: list[str]) -> None:
236
+ missing = [col for col in columns if col not in df.columns]
237
+
238
+ if missing:
239
+ raise ValueError(f"Missing columns: {missing}")